Repository: flatironinstitute/jax-finufft (language: python)

Issue:
batching issue
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!


Discussion / hints:
I released a new version earlier today that should fix this. Want to try `pip install -U jax-finufft` and see if that works? If not, please post a fully functional snippet of code that reproduces this issue. 
It's not working. I get the following now:

```python
/path/to/jax-finufft/lib/python3.7/site-packages/jax_finufft/ops.py in 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     )

    [... skipping hidden 2 frame]

/path/to/jax-finufft/lib/python3.7/site-packages/jax_finufft/ops.py in batch(type_, prim, args, axes, **kwargs)
    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]

/path/to/jax-finufft/lib/python3.7/site-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'
```

The snippet is below. Batching works fine when I replace the call to nufft2 with ordinary matrix multiplication. I have used spacing to isolate the line where I call nufft2 for clarity:

```python
def forward_pass(q_tilde_t, in_array):
    ''' q_tilde_t is a 2 dimensional array '''
    q_hat = jnp.hstack((jnp.zeros(shape=(latent_dim, L_tilde), dtype=np.complex64), q_tilde_t))

    R_nD_nQ_tilde_transpose = jnp.real(nufft2(q_hat, in_array[2, 0:L], eps=1e-6, iflag=-1))

    R_nD_nQ_tilde_conj = R_nD_nQ_tilde_transpose * in_array[1, :L]
    R_nD_nQ_tilde = R_nD_nQ_tilde_transpose.T
    R_nD_nQ_tilde_conj_R_nD_nQ_tilde = jnp.matmul(R_nD_nQ_tilde_conj, R_nD_nQ_tilde)
    R_nD_nQ_tilde_conj_x_kn = jnp.matmul(R_nD_nQ_tilde_conj, in_array[3, 0:L])
    R_nD_nQ_tilde_conj_R_nD_nQ_tilde_dagger = jnp.linalg.pinv(a=R_nD_nQ_tilde_conj_R_nD_nQ_tilde)
    z = jnp.matmul(R_nD_nQ_tilde_conj_R_nD_nQ_tilde_dagger, R_nD_nQ_tilde_conj_x_kn)
    fourier_weights = jnp.matmul(z, q_tilde_t)
    to_return = jnp.hstack((fourier_weights, jnp.array([in_array[0, L]])))

    return to_return

# batch the forward pass 
batch_forward = vmap(forward_pass, in_axes=(None, 0), out_axes=0)
```
Please add all the inputs and imports - I need to be able to execute your code!
```python
import numpy as np
from numpy.random import default_rng
import jax.numpy as jnp
from jax import vmap, random

latent_dim=10 
L_tilde = 60
L = 256
rng = default_rng()

q_tilde_t_r = rng.standard_normal((latent_dim, L_tilde+1))
q_tilde_t_i = rng.standard_normal((latent_dim, L_tilde+1))
q_tilde_t = jnp.array(q_tilde_t_r + 1j * q_tilde_t_i)
in_array = jnp.array(rng.uniform((4, L+1)))
```

As you can see, you can just use random parameters, as long as the shapes are right
sorry, just one thing, in_array is 3 dimensional. so you can set 

```python
in_array = jnp.array(rng.uniform((10, 4, L+1)))
```
I assume you meant: `in_array = jnp.array(rng.uniform(size=(10, 4, L+1)))` (the code you posted doesn't have the expected dimensions).

Can you say more about what you expect to happen here (and perhaps provide a simpler example)? I'm not sure what you expect the behavior to be given these input dimensions.

The behavior that is currently supported by the `vmap` here is:

```python
vmap(nufft2, in_axes=axes, out_axes=axes)(f, x)
# where
f.shape = (..., n_trans, n_k)
x.shape = (..., n_j)
```

where the `...` are the axes we're `vmap`ing over.

I provided that complicated code because I thought you wanted to replicate what I was doing. Basically the data to batch over is In_array, and I want to batch on the 0-th dimension. forward_pass does a bunch of identical matrix multiplications on each slice of in_array. Therefore when I call batch_foward(q_tilde_t, in_array), I expect to apply forward_pass to each slice of in_array. It was critical to provide you with the line batch_foward(q_tilde_t, in_array), sorry for the confusion.
That wasn't really the issue - the issue is that I'm not exactly clear what dimensions you expect everything to have throughout the function. Either way, the current implementation doesn't have support for unbatched arguments, so I'd recommend manually broadcasting your `q_tilde_t` to the right shape and then defining `batch_forward = vmap(forward_pass, in_axes=0, out_axes=0)`, without the `None`. It shouldn't be too hard to implement the proper logic, but I won't have a chance for at least a couple of weeks. I'm happy to keep this issue open as a reminder, or we can open a new one that's clearer about what's needed.
PS the logic will probably look a lot like this, for example:

https://github.com/google/jax/blob/6fa860d5aca1b2cd4323c08394b2d95330328209/jax/_src/lax/linalg.py#L650
Got it thanks 👍 
btw I think you can close this issue, we can open a new one if batching on both doesn't work

Fix patch:
diff --git a/lib/jax_finufft.cc b/lib/jax_finufft.cc
index fb847dd..73f642f 100644
--- a/lib/jax_finufft.cc
+++ b/lib/jax_finufft.cc
@@ -16,17 +16,18 @@ void run_nufft(int type, void *desc_in, T *x, T *y, T *z, std::complex<T> *c, st
   for (int d = 0; d < ndim; ++d) n_k *= descriptor->n_k[d];
 
   nufft_opts *opts = new nufft_opts;
-  typename plan_type<T>::type plan;
   default_opts<T>(opts);
+
+  typename plan_type<T>::type plan;
   makeplan<T>(type, ndim, const_cast<int64_t *>(descriptor->n_k), descriptor->iflag,
               descriptor->n_transf, descriptor->eps, &plan, opts);
   for (int64_t index = 0; index < descriptor->n_tot; ++index) {
-    int64_t j = index * descriptor->n_j * descriptor->n_transf;
+    int64_t i = index * descriptor->n_j;
+    int64_t j = i * descriptor->n_transf;
     int64_t k = index * n_k * descriptor->n_transf;
 
-    setpts<T>(plan, descriptor->n_j, &(x[j]), y_index<ndim, T>(y, j), z_index<ndim, T>(z, j), 0,
+    setpts<T>(plan, descriptor->n_j, &(x[i]), y_index<ndim, T>(y, i), z_index<ndim, T>(z, i), 0,
               NULL, NULL, NULL);
-
     execute<T>(plan, &c[j], &F[k]);
   }
   destroy<T>(plan);
diff --git a/setup.py b/setup.py
index aded446..23dfe72 100644
--- a/setup.py
+++ b/setup.py
@@ -2,6 +2,7 @@
 
 import codecs
 import os
+import shutil
 
 from setuptools import find_packages
 from skbuild import setup
@@ -14,6 +15,11 @@ def read(*parts):
         return f.read()
 
 
+# Otherwise build seems to fail with wrong venv path
+if os.path.exists("_skbuild"):
+    shutil.rmtree("_skbuild")
+
+
 setup(
     name="jax-finufft",
     author="Dan Foreman-Mackey",


Test patch:
diff --git a/tests/ops_test.py b/tests/ops_test.py
index ff3e205..4ba73b0 100644
--- a/tests/ops_test.py
+++ b/tests/ops_test.py
@@ -142,7 +142,6 @@ def test_nufft2_grad(ndim, num_nonnuniform, num_uniform, iflag):
 def test_nufft1_vmap(ndim, num_nonnuniform, num_uniform, iflag):
     random = np.random.default_rng(657)
 
-    eps = 1e-10
     dtype = np.double
     cdtype = np.cdouble
 
@@ -159,7 +158,7 @@ def test_nufft1_vmap(ndim, num_nonnuniform, num_uniform, iflag):
     xs = [jnp.repeat(x_[None], num, axis=0) for x_ in x]
     cs = jnp.repeat(c[None], num, axis=0)
 
-    func = partial(nufft1, num_uniform, eps=eps, iflag=iflag)
+    func = partial(nufft1, num_uniform, iflag=iflag)
     calc = jax.vmap(func)(cs, *xs)
     expect = func(c, *x)
     for n in range(num):
@@ -173,7 +172,6 @@ def test_nufft1_vmap(ndim, num_nonnuniform, num_uniform, iflag):
 def test_nufft2_vmap(ndim, num_nonnuniform, num_uniform, iflag):
     random = np.random.default_rng(657)
 
-    eps = 1e-10
     dtype = np.double
     cdtype = np.cdouble
 
@@ -187,11 +185,28 @@ def test_nufft2_vmap(ndim, num_nonnuniform, num_uniform, iflag):
     f = f.astype(cdtype)
 
     num = 5
-    xs = [jnp.repeat(x_[None], num, axis=0) for x_ in x]
-    fs = jnp.repeat(f[None], num, axis=0)
+    xs = [jnp.repeat(x_[jnp.newaxis], num, axis=0) for x_ in x]
+    fs = jnp.repeat(f[jnp.newaxis], num, axis=0)
 
-    func = partial(nufft2, eps=eps, iflag=iflag)
+    func = partial(nufft2, iflag=iflag)
     calc = jax.vmap(func)(fs, *xs)
     expect = func(f, *x)
     for n in range(num):
         np.testing.assert_allclose(calc[n], expect)
+
+
+def test_multi_transform():
+    random = np.random.default_rng(314)
+
+    n_tot, n_tr, n_j, n_k = 4, 10, 100, 12
+    f_shape = (n_tot, n_tr, n_k)
+    c_shape = (n_tot, n_tr, n_j)
+    f = random.standard_normal(size=f_shape) + 1j * random.standard_normal(size=f_shape)
+    c = random.standard_normal(size=c_shape) + 1j * random.standard_normal(size=c_shape)
+    x = random.uniform(-np.pi, np.pi, (n_tot, n_j))
+
+    calc1 = nufft1(n_k, c, x)
+    calc2 = nufft2(f, x)
+    for n in range(n_tr):
+        np.testing.assert_allclose(calc1[:, n], nufft1(n_k, c[:, n], x))
+        np.testing.assert_allclose(calc2[:, n], nufft2(f[:, n], x))