Hi,

I get the following error when I try to batch nufft2 in Jax.

process_primitive(self, primitive, tracers, params) 161 frame = self.get_frame(vals_in, dims_in) 162 batched_primitive = self.get_primitive_batcher(primitive, frame) --> 163 val_out, dim_out = batched_primitive(vals_in, dims_in, **params) 164 if primitive.multiple_results: 165 return map(partial(BatchTracer, self), val_out, dim_out)

TypeError: batch() got an unexpected keyword argument 'output_shape'

it seems like this is caused by

nufft2(source, iflag, eps, *points) 57 58 return jnp.reshape( ---> 59 nufft2_p.bind(source, *points, output_shape=None, iflag=iflag, eps=eps), 60 expected_output_shape, 61 )

Is there something I am doing wrong?

Thanks for your help!

Hi,

I am posting this here to give a clearer description of the problem that I am having.

When I run the following snippet, I would like to have A be a 4 by 100 array whose I-th row is the output of linear_func(q, X[I,:]).

import numpy as np
import jax.numpy as jnp
from jax import vmap
from jax_finufft import nufft2

rng = np.random.default_rng(seed=314)

d=10 
L_tilde = 10
L = 100

qr = rng.standard_normal((d, L_tilde+1))
qi = rng.standard_normal((d, L_tilde+1))
q = jnp.array(qr + 1j * qi)
X = jnp.array(rng.uniform(low=0.0, high=1.0, size=(4, L)))

def linear_func(q, x):
  v = jnp.ones(shape=(1, L_tilde))

  return jnp.matmul(v, nufft2(q, x, eps=1e-6, iflag=-1))

batched = vmap(linear_func, in_axes=(None, 0), out_axes=0)

A = batched(q, X)

However, when I run the snippet I get the error posted below. You had said last time that there could be a work around for unbatched arguments but I could not figure it out.

Here is the error:

UnfilteredStackTrace                      Traceback (most recent call last)
<ipython-input-25-c23bc4fa7eb4> in <module>()
----> 1 test = batched(q, X)

30 frames
UnfilteredStackTrace: TypeError: '<' not supported between instances of 'NoneType' and 'int'

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

TypeError                                 Traceback (most recent call last)
/usr/local/lib/python3.7/dist-packages/jax_finufft/ops.py in <genexpr>(.0)
    281     else:
    282         mx = args[0].ndim - ndim - 1
--> 283     assert all(a < mx for a in axes)
    284     assert all(a == axes[0] for a in axes[1:])
    285     return prim.bind(*args, **kwargs), axes[0]

TypeError: '<' not supported between instances of 'NoneType' and 'int'

I hope this gives a clearer picture than what I had last time. Thanks so much for your help!

I am having trouble installing jax-finuff even with the instructions on the home page. The installation fails with

ERROR: Failed building wheel for jax-finufft Failed to build jax-finufft ERROR: Could not build wheels for jax-finufft, which is required to install pyproject.toml-based projects

Hi Dan,

First, thank you for releasing this package, I was very glad to find it!

I noted the error below when attempting to differentiate nufft1 with respect to points. I am very new to JAX so I could be mistaken, but I don't believe this is the intended behavior:

from jax_finufft import nufft1, nufft2
import numpy as np
import jax.numpy as jnp
from jax import grad
M = 100000
N = 200000

x = 2 * np.pi * np.random.uniform(size=M)
c = np.random.standard_normal(size=M) + 1j * np.random.standard_normal(size=M)

def norm_nufft1(c,x):
    f = nufft1(N, c, x, eps=1e-6, iflag=1)
    return jnp.linalg.norm(f)

def norm_nufft2(c,x):
    f = nufft2( c, x, eps=1e-6, iflag=1)
    return jnp.linalg.norm(f)

grad(norm_nufft2,argnums =(1))(c,x) # Works fine
grad(norm_nufft1,argnums =(0,))(c,x) # Works fine
grad(norm_nufft1,argnums =(0,1))(c,x) # Works fine
grad(norm_nufft1,argnums =(1))(c,x) # Throws error

The error is below:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/var/folders/dg/zzj57d7d1gs1k9l8sdfh043m0000gn/T/ipykernel_13430/3134504975.py in <module>
     22 grad(norm_nufft1,argnums =(0,))(c,x) # Works fine
     23 grad(norm_nufft1,argnums =(0,1))(c,x) # Works fine
---> 24 grad(norm_nufft1,argnums =(1))(c,x) # Throws error
     25 
     26 

    [... skipping hidden 10 frame]

/var/folders/dg/zzj57d7d1gs1k9l8sdfh043m0000gn/T/ipykernel_13430/3134504975.py in norm_nufft1(c, x)
     11 
     12 def norm_nufft1(c,x):
---> 13     f = nufft1(N, c, x, eps=1e-6, iflag=1)
     14     return jnp.linalg.norm(f)
     15 

    [... skipping hidden 21 frame]

~/opt/anaconda3/envs/alphafold/lib/python3.8/site-packages/jax_finufft/ops.py in nufft1(output_shape, source, iflag, eps, *points)
     39 
     40     return jnp.reshape(
---> 41         nufft1_p.bind(source, *points, output_shape=output_shape, iflag=iflag, eps=eps),
     42         expected_output_shape,
     43     )

    [... skipping hidden 3 frame]

~/opt/anaconda3/envs/alphafold/lib/python3.8/site-packages/jax_finufft/ops.py in jvp(type_, prim, args, tangents, output_shape, iflag, eps)
    248 
    249         axis = -2 if type_ == 2 else -ndim - 1
--> 250         output_tangent *= jnp.concatenate(jnp.broadcast_arrays(*scales), axis=axis)
    251 
    252         expand_shape = (

~/opt/anaconda3/envs/alphafold/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py in concatenate(arrays, axis)
   3405   if hasattr(arrays[0], "concatenate"):
   3406     return arrays[0].concatenate(arrays[1:], axis)
-> 3407   axis = _canonicalize_axis(axis, ndim(arrays[0]))
   3408   arrays = _promote_dtypes(*arrays)
   3409   # lax.concatenate can be slow to compile for wide concatenations, so form a

~/opt/anaconda3/envs/alphafold/lib/python3.8/site-packages/jax/_src/util.py in canonicalize_axis(axis, num_dims)
    277   axis = operator.index(axis)
    278   if not -num_dims <= axis < num_dims:
--> 279     raise ValueError(f"axis {axis} is out of bounds for array of dimension {num_dims}")
    280   if axis < 0:
    281     axis = axis + num_dims

ValueError: axis -2 is out of bounds for array of dimension 1

If none of the points are batched in a vmap, we should be able to get a faster computation by stacking the transforms into a single transform and then reshaping. It might be worth making the stacked axes into an explicit parameter rather than just trying to infer it, but that would take some work on the interface.

In most cases, we'll just need to broadcast all the inputs out to the right shapes (this shouldn't be too hard), but when none of the points get mapped, we can get a bit of a speed up by stacking the transforms. This starts to implement that logic, but it's not quite ready yet.

@lgarrison and I have been chatting about the possibility of writing a paper describing what we're doing here. Something like "Differentiable programming with NUFFTs" or "NUFFTs for machine learning applications" or something. This issue is here to remind me to look into possible venues for this.

So far I just have the CMake definitions to compile cuFINUFFT when nvcc is found, but I haven't started writing the boilerplate needed to loop it into XLA. Coming soon!

Keeping @lgarrison in the loop.

There aren't currently any plans to support the Type 3 transform for a few reasons:

• I'm not totally sure of the use cases,
• The logic will probably be somewhat more complicated than the existing implementations, and
• cuFINUFFT doesn't seem to support Type 3.

Do we want to work around these issues?

Currently, if any of the finufft methods fail with a non-zero error we just ignore it and keep on trucking. How should we handle this? It looks like the JAX convention is currently to set everything to NaN when an op fails, since propagating errors up from XLA can be a bit of a pain.

Via cufinufft. We'll need to figure out how to get XLA's handling of CUDA streams to play nice with cufinufft (this is way above my pay grade). Some references:

1. A simple example of how to write an XLA compatible CUDA kernel: https://github.com/dfm/extending-jax/blob/main/lib/kernels.cc.cu
2. The source code for how JAX wraps cuBLAS: https://github.com/google/jax/blob/main/jaxlib/cublas_kernels.cc
3. There is a tensorflow implementation that might have some useful context: https://github.com/mrphys/tensorflow-nufft/blob/master/tensorflow_nufft/cc/kernels/nufft_kernels.cu.cc

Perhaps @lgarrison is interested :D