Neojax

neojax (Neural Operators in JAX) is an implementation of Neural Operators built on top of JAX and Equinox. It provides a clean, modular API inspired by the original neuraloperator library.

It is designed to be fully compatible with all JAX features such as vmap, jit, and grad.

One of several model architecture neojax implements is...

GeoFNO: 2D wave propagation with spatially varying speed

Trained (partially) in docs/examples/05_advanced_training.ipynb — predicts wave fields over heterogeneous media with a single forward pass.

Available Features

Currently, neojax is in its early stages. Expect possible breaking changes.

Only the following models and features are available:

Installation

Install the python package via pypi

pip3 install neojax-operators

The core library is designed to have as few dependencies as possible. Some submodules therefore rely on additional dependencies. Install them as needed:

• Running the examples: neojax-operators[ex]
• Using the PDE data generation: neojax-operators[gen]
• Using data downloading and some other data features: neojax-operators[data]
• Using the benchmark module: neojax-operators[benchmark]

Quickstart

neojax exposes a similar API to neuraloperators and equinox and should therefore be familiar to use:

import jax.numpy as jnp
import jax.random as jr
from neojax.models import FNO

key = jr.key(0)

fno = FNO(
      key=key,
      modes=(12, 12),
      hidden_channels=64,
      in_channels=2,
      out_channels=1,
      n_layers=2
)

x = jnp.ones((2, 64, 64))

pred = fno(x)

For a more detailed introduction refer to the examples in the documentation.

Performance

neojax leverages JAX's static compilation to outperform PyTorch, especially for spectral operators like FNO.

Benchmark plot (1D Burgers, NVIDIA A100)

Full methodology, more PDEs, and reproduction scripts → benchmarks/

Motivation

JAX is widely used in Scientific Machine Learning (SciML) for its functional transformations (jit, vmap, grad) and NumPy-like API. Neural Operators are a common approach for solving PDEs in this space, but a native JAX implementation was missing. neojax fills that gap.

Design Choices

neojax started as a port of the PyTorch neuraloperator library, but was rewritten from the ground up as a JAX-native implementation instead. Porting PyTorch idioms directly into a functional framework added unnecessary complexity; building on Equinox gives neojax a class-based API while staying consistent with JAX's pure-functional design.

Roadplan

In upcoming releases more models and components will be added in roughly the following order:

1. LocalNO, SFNO, RNO, and others
2. Debugging utilities

And much more to come!

Contributions

If you'd like to contribute any features, models, or fix implementation errors, please do so. Any contributions are appreciated. Have a look at the CONTRIBUTING.md guide for details on how to do so. I am also open to advice on restructuring and any other design choices that could be improved.

Citation

If you use neojax in your research, please cite it using the following BibTeX entry:

@software{neojax,
  author = {Paul Gekeler},
  title = {neojax: Neural Operators in Jax},
  year = {2026},
  url = {https://github.com/paulgekeler/neojax}
}