nufftax 0.2.0

Pure JAX implementation of Non-Uniform FFT

nufftax

Pure JAX implementation of the Non-Uniform Fast Fourier Transform (NUFFT)

---

Why nufftax?

A JAX package for NUFFT already exists: jax-finufft. However, it wraps the C++ FINUFFT library via Foreign Function Interface (FFI), exposing it through custom XLA calls. This approach can lead to:

• Kernel fusion issues on GPU — custom XLA calls act as optimization barriers, preventing XLA from fusing operations
• CUDA version matching — GPU support requires matching CUDA versions between JAX and the library

nufftax takes a different approach — pure JAX implementation:

• Fully differentiable — gradients w.r.t. both values and sample locations
• Pure JAX — works with jit, grad, vmap, jvp, vjp with no FFI barriers
• GPU ready — runs on CPU/GPU without code changes, benefits from XLA fusion
• All NUFFT types — Type 1, 2, 3 in 1D, 2D, 3D

JAX Transformation Support

Transform	jit	grad/vjp	jvp	vmap
Type 1 (1D/2D/3D)	✅	✅	✅	✅
Type 2 (1D/2D/3D)	✅	✅	✅	✅
Type 3 (1D/2D/3D)	✅	⚠️	⚠️	✅

✅ = Fully supported | ⚠️ = Work in progress

Differentiable inputs:

• Type 1: grad w.r.t. c (strengths) and x, y, z (coordinates)
• Type 2: grad w.r.t. f (Fourier modes) and x, y, z (coordinates)
• Type 3: grad w.r.t. c and coordinates (in development)

Installation

uv pip install nufftax

Quick Example

import jax
import jax.numpy as jnp
from nufftax import nufft1d1

# Irregular sample locations in [-pi, pi)
x = jnp.array([0.1, 0.7, 1.3, 2.1, -0.5])
c = jnp.array([1.0+0.5j, 0.3-0.2j, 0.8+0.1j, 0.2+0.4j, 0.5-0.3j])

# Compute Fourier modes
f = nufft1d1(x, c, n_modes=32, eps=1e-6)

# Differentiate through the transform
grad_c = jax.grad(lambda c: jnp.sum(jnp.abs(nufft1d1(x, c, n_modes=32)) ** 2))(c)

Documentation

Read the full documentation →

• Quickstart — get running in 5 minutes
• Concepts — understand the mathematics
• Tutorials — MRI reconstruction, spectral analysis, optimization
• API Reference — complete function reference

License

MIT. Algorithm based on FINUFFT by the Flatiron Institute.