python实现可在GPU运行修正的第一类贝塞尔函数

Python的Scipy库里边是有修正的贝塞尔函数scipy.special.i0的，但是它只是一个数值计算的过程。在深度学习需要用到贝塞尔函数的时候，显然scipy.special.i0并不能对其求梯度，因此本人根据网上资料手撸了一个可对其求梯度的修正的贝塞尔函数，同时它可在CPU或是GPU上运行。

修正的第一类贝塞尔函数(CPU版)

如下是修正的第一类贝塞尔函数(0阶)

// Bessel_i0() 相当于 scipy.special.i0() 
def Bessel_i0(a):    //ver7.0 modify bessel function  
    temp = torch.zeros(a.shape).type(torch.float64) 
    
    y = a[abs(a)<3.75]/3.75 
    y = y**2
    ans = 1.0 + y * (3.5156229 + y * (3.0899424 + y * (1.2067492 + y * (0.2659732 + y * (0.360768e-1 + y * 0.45813e-2)))))
    temp[abs(a)<3.75] = ans.type(torch.float64)  
   
    y = 3.75/a[abs(a)>=3.75]
    y.type(dtype)
    ans = (torch.exp(a[abs(a)>=3.75]) / torch.sqrt(a[abs(a)>=3.75]))*(0.39894228 + y * (0.1328592e-1 + y * (0.225319e-2 + y * (-0.157565e-2 + y * (0.916281e-2 + y * (-0.2057706e-1 + y * (0.2635537e-1 + y * (-0.1647633e-1 + y * 0.392377e-2)))))))) 
    temp[abs(a)>=3.75] = ans.type(torch.float64)  
    
    return temp

如下是修正的第一类贝塞尔函数(1阶)