Custom bessel autograd function

alexmm: <p>Hi, I was able to implement a class for Jn to have arbitrary derivatives as shown below:</p>
<pre><code class="lang-auto">import numpy as np
from scipy.special import jv
import torch


class besselJv(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, v):
        ctx.save_for_backward(x, v)
        return jv(v, x)

    @staticmethod
    def backward(ctx, grad_out):
        x, v = ctx.saved_tensors
        
        return grad_out*0.5*(besselJv.apply(x, v-1) -  besselJv.apply(x, v+1)), None


x = torch.tensor([2.0], dtype=torch.double, requires_grad=True)
v = torch.tensor(1)
y = besselJv.apply(x, v)

dy_dx   = torch.autograd.grad(y, x, create_graph=True)
dy2_dx2 = torch.autograd.grad(dy_dx, x)

print(dy_dx[0].item())
print(dy2_dx2[0])
</code></pre>
<p>However when I try implementing this using <code>v</code> as a constant and not a tensor following this post:</p><aside class="quote" data-post="1" data-topic="18609">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/i/b5ac83/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/modified-bessel-function-of-order-0/18609">Modified Bessel Function of Order 0</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    How to implement modified bessel functions in Pytorch? Any ideas/hints?
  </blockquote>
</aside>

<pre><code class="lang-auto">import numpy as np
from scipy.special import jv
import torch


class besselJv(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, v):
        ctx.save_for_backward(x)
        ctx._v = v
        return torch.from_numpy(jv(v, x.detach().numpy()))

    @staticmethod
    def backward(ctx, grad_out):
        x = ctx.saved_tensors
        v = ctx._v
        return grad_out*0.5*(besselJv.apply(x, v-1) -  besselJv.apply(x, v+1)), None


x = torch.tensor([2.0], dtype=torch.double, requires_grad=True)
v = torch.tensor(1)
y = besselJv.apply(x, v)

dy_dx   = torch.autograd.grad(y, x, create_graph=True)
dy2_dx2 = torch.autograd.grad(dy_dx, x)

print(dy_dx[0].item())
print(dy2_dx2[0])
</code></pre>
<p>I get an error with this line:</p>
<p>What is the issue here? Is this code outdated?</p>
<p>Thank you,<br>
Alex</p>

ptrblck: <p>I don’t know which error you are seeing, as it seems to be missing from your post, but executing your code yields:</p>
<pre><code class="lang-python">TypeError: 'besselJv' object is not iterable
</code></pre>
<p>for me.<br>
Checing the output <code>y</code> also shows:</p>
<pre><code class="lang-python">y
&lt;__main__.besselJv at 0x7fbfac055b80&gt;
</code></pre>
<p>which is an object and not a tensor as expected so I guess you want to use:</p>
<pre><code class="lang-python">y = besselJv.apply(x, v)
</code></pre>
<p>instead which works fine for me.</p>

alexmm: <p>Hi thank you for the prompt reply. Yes the first code snippet works, but I would like to have <code>v</code> not as a tensor but as a constant like in the second code snippet, as was done in the referenced post.</p>

ptrblck: <p>I’m not sure what you are looking exactly for, but I assume you want to save <code>v</code> in <code>ctx._v</code>?<br>
If so, then your second code snippet contains a few slight issues and this should work:</p>
<pre><code class="lang-python">class besselJv(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, v):
        ctx.save_for_backward(x)
        ctx._v = v
        return torch.from_numpy(jv(v.numpy(), x.detach().numpy()))

    @staticmethod
    def backward(ctx, grad_out):
        x, = ctx.saved_tensors
        v = ctx._v
        return grad_out*0.5*(besselJv.apply(x, v-1) -  besselJv.apply(x, v+1)), None


x = torch.tensor([2.0], dtype=torch.double, requires_grad=True)
v = torch.tensor(1)
y = besselJv.apply(x, v)

dy_dx   = torch.autograd.grad(y, x, create_graph=True)
dy2_dx2 = torch.autograd.grad(dy_dx, x)

print(dy_dx[0].item())
print(dy2_dx2[0])
</code></pre>

alexmm: <p>Great, that works! You can simplify it further by using <code>v=1</code> instead of <code>v = torch.tensor(1)</code>, and replacing <code>v.numpy()</code> with <code>v</code> in the return of the forward function.</p>