Enabling torch.Tensor datatype support on CPU device

[remy-abergel]

Hi,

The following code works fine when device='cuda' and fails when device='cpu' :

##################
# import modules #
##################
import cufinufft 
import finufft
import torch 

#########################################
# select torch device ('cpu' or 'cuda') #
#########################################
device = 'cuda' # try also device = 'cpu'
lib = finufft if device == 'cpu' else cufinufft

####################################################
# draw M nonuniform points in the square [0,2pi)^2 #
####################################################
M = 100000
x = 2 * torch.pi * torch.rand(M).to(device)
y = 2 * torch.pi * torch.rand(M).to(device)

###################################
# sample random complex strengths #
###################################
c = (torch.rand(M) + 1J*torch.rand(M)).to(device)

####################################################################
# calculate the transform over the desired number of Fourier modes #
####################################################################
N1 = 1000
N2 = 2000
f = lib.nufft2d1(x, y, c, (N1,N2))

When device='cpu', the following error is raised:

RuntimeError: FINUFFT x dtype should be float64 for double precision or float32 for single precision

A trivial workaround can be made using

f = torch.from_numpy(lib.nufft2d1(x.numpy(), y.numpy(), c.numpy(), (N1,N2)))

Both finufft and cufinufft packages are probably not intended to handle torch.Tensor arrays on CPU device, but adding such support into one of those packages would be nice to allow device independent development in PyTorch. What do you think?

[lu1and10]

There is a much nicer PyTorch package https://github.com/flatironinstitute/pytorch-finufft by @WardBrian and @eickenberg which calls cufinufft and finufft as backend. You may consider use pytorch-finufft directly, it also has autodiff.

[remy-abergel]

Thank you for this input, I will take a look at this wrapper.

Still, I am wondering if the error (caused by these tests in the source code) could be ignored and if the rest of the code is able to deal with torch.Tensor arrays "as is", using numpy interoperability. I will take a look a this as well.

[ahbarnett]

@janden could comment on this. Is it a matter of switching off the tests, to include objects that are not numpy arrays?

[janden]

At this point, FINUFFT does not support PyTorch tensors as input. I think the natural solution is to point users toward pytorch-finufft since this will pick the appropriate library depending on the device and includes autodiff, as @lu1and10 points out.

That being said, maybe we should discuss this. I'm not sure what would be involved in the Python code to support this. My guess is that it would be somewhat more involved than just relying on NumPy interoperability, but not too much. On the other hand, not having autodiff would limit its usefulness.