Implementation of scipy function as Op

[AuguB]

Hi all, I am working on a custom Op that I have to use in a project, and I may want to submit a PR for this function (the scipy.special.kv function) and perhaps others in that module if it goes well, but I want to make sure that I am doing it correctly first.

Here's my Op as it is now:

from scipy.special import kv
class K(Op):
    """
    Modified Bessel function of the second kind, pytensor implementation
    """
    __props__ = ()

    def make_node(self, p, x):
        p = pt.tensor.as_tensor_variable(p)
        x = pt.tensor.as_tensor_variable(x)
        return Apply(self, [p, x], [p.type()])

    def perform(self, node, inputs_storage, output_storage):
        output_storage[0][0] = kv(inputs_storage[0], inputs_storage[1])

    def grad(self, inputs, output_grads):
        dp = 1e-16
        p = inputs[0]
        x = inputs[1]
        grad = (self(p + dp, x) - self(p - dp, x)) / (2*dp)
        return [output_grads[0] * grad, grad_not_implemented('KOp', 1, 'x', 'Grad not implemented for KOp')]

Specifically, I need the gradient with respect to the first parameter, p, and currently I am using a numerical approximation, because there exists no analytical expression for it (afaik). The performance of this Op in a PyMC setting is not so good, so I would like to speed it up if possible.

My questions are the following:

1. Is the numerical approximation appropriate, and if not, are there other ways to do this such that I can use the nutpie sampler (which fails if I remove the grad function).
2. When using this Op in a PyMC setting and sampling with nutpie, I get the warning that numba will use object mode for KOp. Are there ways to make numba work with this scipy special function? I have tried to follow the guide here: https://pytensor.readthedocs.io/en/latest/extending/creating_a_numba_jax_op.html, but I have had limited success. I have also tried to use the numba-scipy package, but also there I have not been successful.

Any help is appreciated!

[ricardoV94]

Hi @AuguB

An Op like this would be implemented as a ScalarOp (see pytensor.scalar.math.py) and then "vectorized" via Elemwise.

It's fine for the gradients to use numerical approximation. We do this for some like GammaInc functions, you can find it in the same file.

In this case scipy also seems to have a helper to compute it: https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.kvp.html#scipy.special.kvp

We can try to wrap it as it's own Op or implement directly in PyTensor for support in other backends

For numba support of kv you'll need a dispatch for the ScalarOp. There's not much we can do there automatically. Some cases we try to access Scipy C functions directly. You can try to see if the same approach would work for this Op.

Contributions are welcome of course!

[ricardoV94]

I see now that you were suggesting using finite-differences. That is probably not good enough for us (but may be fine as your custom Op). By the way the gradient wrt to x is pretty simple, just (-v / x) * kv(v, x) - kv(v-1, x).

For the gradient wrt to v, I found this paper and accompanying JAX/Tfp code that you may find interesting: https://github.com/tk2lab/logbesselk/tree/main

It's for the log of Kv, not Kv directly, but it's straightforward to translate if you are not interested in the log representation itself.

I'm opening a PR to add the Kv and gradient wrt to x, as it showed up in another project I was working on. We can chat about expanding it to cover the gradient wrt to v as well.

[AuguB]

Hi @ricardoV94, thanks for the replies, they have been very helpful.

I've currently implemented the Kv (and its derivative) as BinaryScalarOps. And I'm happy with how it looks and runs for my case for now (although speedups are always welcome).

Having Kv and the gradient wrt v available in pymc sounds great, but I don't think I understand the different backends well enough to know what the best (and most general) solution is. If I understand correctly, the repo you linked provides solutions for Jax and TF, but not for Numba or pure Python. I have looked to providing C code for the K op directly, because it is fast and portable, but I am not sure if Numba, Jax, and TF support comes automatically (and my experience with C is meager at best). I would be up for working on this, what do you think is the best way to proceed?

[ricardoV94]

That JAX implementation can be written almost verbatim in PyTensor, which means all backends would be automatically supported.