Error: 'expected np.ndarray (got numpy.float64)' when using torch.sum()

autograd
1

Hi, I am trying to write functions like p_spatial_comp_real to get their gradient as shown below.
However when I try running this code I get this error: TypeError: expected np.ndarray (got numpy.float64). This seems to be related to using the function torch.sum() that turns the tensor array into a scalar. Is there a way of doing this sum and outputting a 1 by 1 tensor array? I tested the function modJv using x = torch.tensor([2.4]) and it works.

Thank you,
Alex

import numpy as np
from scipy.special import jv
import torch

#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()
    
    sinTheta2 = torch.cross(distVec,transNormal).pow(2).sum()/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)

#Parts of sound field that are spatialy dependent and don't have complex exponents
p_spatial_comp, tpDist = pressureSpatialCompFunc(point, transLocation,transNormal, k*r_trans)

#Compute real and imaginary parts of spatialy depedent component of transducer sound field
p_spatial_comp_real = p_spatial_comp*torch.sin(k*tpDist)
p_spatial_comp_imag = p_spatial_comp*torch.cos(k*tpDist)
2

There should be a variety of ways to achieve the desired shape e.g.,

>>> a = torch.ones(10, 10)
>>> torch.sum(a, (0, 1), keepdim=True)
tensor([[100.]])
>>> torch.sum(a).unsqueeze(dim=0)
tensor([100.])
3

Thank you eqy! That works. Do you see any way not to need this by changing other types in the code? Is this the simplest solution?