AttributeError: 'float' object has no attribute 'dtype' when using extra arguments in jacrev()

alexmm: <p>Hi, I’m getting an error due to passing an argument (I think its the float k*r_trans) when using extra arguments to the function <code>jacrev()</code>. I tried declaring this as a tensor and then converting it back to a float inside the function <code>pressureSpatialCompRealFunc</code> but then it says the new variable doesnt have storage. Is there a way of passing this variable? The code is below.</p>
<p>Also, when using <code>jacfwd()</code> on the output of <code>jacrev()</code>, do I need the extra parameters?</p>
<p>Thank you,<br>
Alex</p>
<pre><code class="lang-auto">import numpy as np
from scipy.special import jv
import torch

import time


#Autograd class declaration of function J_v(sqrt(x)) for pytorch
#This implementation does not have the machine precision error due to /sqrt(x) when x is or is close to 0
class modJv(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, v, a):
        ctx.save_for_backward(x)
        ctx._v = v
        ctx._a = a
        return torch.from_numpy(jv(v, np.array(a*np.sqrt(x.detach().numpy()))))

    @staticmethod
    def backward(ctx, grad_out):
        x, = ctx.saved_tensors
        v = ctx._v
        a = ctx._a
        
        Jv   = modJv.apply(x,   v,a)
        Jvm2 = modJv.apply(x, v-2,a)
        Jv2  = modJv.apply(x, v+2,a)
        return grad_out*a*a/8.0*((Jv+Jvm2)/(v-1) - (Jv2+Jv)/(v+1)), None, None


#Function to compute parts of sound field that are spatialy dependent and don't have complex exponents
def pressureSpatialCompFunc(point, transLocation,transNormal, kr_trans):
    dirV = 1 #Nu parameter for Bessel function in original form of transducer field directivity function
    
    distVec   = point-transLocation
    dist2     = distVec.pow(2).sum().unsqueeze(dim=0)
    sinTheta2 = torch.cross(distVec,transNormal).pow(2).sum().unsqueeze(dim=0)/dist2
    
    directivity = 0.25*(modJv.apply(sinTheta2,dirV+1,kr_trans) + modJv.apply(sinTheta2,dirV-1,kr_trans))
    
    dist = dist2.sqrt()
    return directivity/dist, dist



#All values in SI units
c_m     = 332.0  #Speed of sound in air (m/s)
f_trans = 40.0e3 #Transducer frequency (40 kHz)

k = 2*np.pi*f_trans/c_m #Sound wavenumber

r_trans = 5.0e-3 #Transducer radius (5 mm)


#Transducer location
transLocation = torch.tensor([4.0e-3, 2.0e-3,  1.0e-3], dtype=torch.double, requires_grad=False)
#Transducer unit normal vector
transNormal   = torch.tensor([0.0,    0.0,     1.0   ], dtype=torch.double, requires_grad=False)
#Point where sound field is measured
point         = torch.tensor([4.0e-3, 2.0e-3, 20.0e-3], dtype=torch.double, requires_grad=True)


#Compute real and imaginary parts of spatialy depedent component of transducer sound field
def pressureSpatialCompRealFunc(point, transLocation,transNormal, kr_trans):
    p_spatial_comp, tpDist = pressureSpatialCompFunc(point, transLocation,transNormal, kr_trans)
    p_spatial_comp_real = p_spatial_comp*torch.sin(k*tpDist)


start_time = time.time()
#p_spatial_comp, tpDist = pressureSpatialCompFunc(point, transLocation,transNormal, kr_trans)
#p_spatial_comp_real = p_spatial_comp*torch.sin(k*tpDist)
#jac_f = torch.autograd.grad(p_spatial_comp_real, point, create_graph=True)
jac_f = torch.func.jacrev(pressureSpatialCompRealFunc,argnums=3)(point, transLocation,transNormal,k*r_trans)
jac_time   = time.time() - start_time
print("--- %s seconds ---" % jac_time)

start_time = time.time()
hess_f     = torch.func.jacfwd(jac_f,argnums=0)(point)
hess_time  = time.time() - start_time
print("--- %s seconds ---" % hess_time)
print("--- %s seconds ---" % (jac_time + hess_time))


print(jac_f)
print(hess_f)
</code></pre>

ptrblck: <aside class="quote no-group" data-username="alexmm" data-post="1" data-topic="176331">
<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/a/b9bd4f/48.png" class="avatar"> alexmm:</div>
<blockquote>
<p>I tried declaring this as a tensor and then converting it back to a float inside the function <code>pressureSpatialCompRealFunc</code> but then it says the new variable doesnt have storage.</p>
</blockquote>
</aside>
<p>That’s strange, as I see this error:</p>
<pre data-code-wrap="python"><code class="lang-python">RuntimeError: In order to use an autograd.Function with functorch transforms (vmap, grad, jvp, jacrev, ...), it must override the setup_context staticmethod. For more details, please see https://pytorch.org/docs/master/notes/extending.func.html
</code></pre>
<p>after passing <code>torch.tensor(k*r_trans)</code> to the call, which points to <a href="https://pytorch.org/docs/master/notes/extending.func.html">these docs</a>.<br>
Are you seeing an internal error instead?</p>

alexmm: <p>Hi ptrblock, when I try passing the variable like you did, as shown below, I get the following error:<br>
<code>RuntimeError: Cannot access data pointer of Tensor that doesn't have storage</code></p>
<pre><code class="lang-auto">def pressureSpatialCompRealFunc(point, transLocation,transNormal, kr_trans):
    kr_trans = kr_trans.detach().numpy()
    p_spatial_comp, tpDist = pressureSpatialCompFunc(point, transLocation,transNormal, kr_trans)
    p_spatial_comp_real = p_spatial_comp*torch.sin(k*tpDist)

kr_trans = torch.tensor(k*r_trans)
start_time = time.time()
jac_f = torch.func.jacrev(pressureSpatialCompRealFunc,argnums=3)(point, transLocation,transNormal,kr_trans)
jac_time   = time.time() - start_time
print("--- %s seconds ---" % jac_time)
</code></pre>

soulitzer: <p>You can’t call .numpy() on tensors (outside of <code>forward</code> of custom autograd Function) when performing using torch.func transforms like jacrev. If you cannot avoid that numpy call should move it  to be done inside custom Function.</p>
<p>(this is expected, but the error definitely needs to be improved)</p>

alexmm: <p>Actually I’ve been using <code>argnums</code> wrong; from the documentation in jacrev, the code below should work. I want to do the jacobian w.r.t <code>point</code>, and I want the function value as well, so I’m using <code>has_aux=True</code>. However when I run this now I get this error: <code>AttributeError: 'float' object has no attribute 'dtype'</code></p>
<pre><code class="lang-auto">c_m     = 332.0
f_trans = 40.0e3
k       = 2*np.pi*f_trans/c_m 
r_trans = 5.0e-3

transLocation = torch.tensor([4.0e-3, 2.0e-3,  1.0e-3], dtype=torch.double, requires_grad=False)
transNormal   = torch.tensor([0.0,    0.0,     1.0   ], dtype=torch.double, requires_grad=False)
point         = torch.tensor([4.0e-3, 2.0e-3, 20.0e-3], dtype=torch.double, requires_grad=True)

def pressureSpatialCompRealFunc(point, transLocation, transNormal, kr_trans):
    p_spatial_comp, tpDist = pressureSpatialCompFunc(point, transLocation, transNormal, kr_trans)
    p_spatial_comp_real = p_spatial_comp*torch.sin(k*tpDist)

jac_f, f_x = torch.func.jacrev(pressureSpatialCompRealFunc, argnums=0, has_aux=True) \
(point, transLocation, transNormal, k*r_trans)
</code></pre>