norax 1.0.0

Neural operators in JAX

norax

Neural operators in JAX.

norax provides implementations of neural operators built on top of JAX and Equinox.

Neural operators learn mappings between function spaces.

Currently only Fourier neural operators (Li et al., 2020) are provided in this library.

Check out the documentation page for more details.

Installation

pip install norax

Or with uv:

uv add norax

Quick start

import jax
import jax.numpy as jnp
from norax.models import FNO

key = jax.random.key(0)

# Define a 1D Fourier neural operator
model = FNO(
    key,
    channels_in=2,   # 1D scalar field: [x-coordinate, field value]
    channels_out=1,
    n_modes=(16,),   # Fourier modes to retain per spatial axis
    width=64,
    depth=4,
)

# Forward pass over a batch using vmap
x = jnp.ones((10, 64, 2))  # (batch, *grid_shape, channels_in)
y = jax.vmap(model)(x)     # (batch, *grid_shape, channels_out)

The n_modes tuple determines the spatial dimensionality of the operator: a 1-tuple gives a 1D FNO, a 2-tuple gives a 2D FNO, and so on.

Contributing

See CONTRIBUTING.md for instructions on setting up a development environment, running tests, linting, and building the documentation locally.