Metadata-Version: 2.4
Name: neojax-operators
Version: 0.0.5
Summary: Neural Operators in JAX
Project-URL: Homepage, https://github.com/paulgekeler/neojax
Author-email: Paul Gekeler <pgekeler@gmail.com>
License: MIT
License-File: LICENSE
Keywords: deep-learning,fno,jax,neural-operators,scientific-computing
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Requires-Dist: equinox
Requires-Dist: jax
Requires-Dist: jaxlib
Provides-Extra: dev
Requires-Dist: mypy; extra == 'dev'
Requires-Dist: neuraloperator; extra == 'dev'
Requires-Dist: pytest>=8.0.0; extra == 'dev'
Requires-Dist: ruff; extra == 'dev'
Requires-Dist: torch>=2.0.0; extra == 'dev'
Provides-Extra: docs
Requires-Dist: mkdocs; extra == 'docs'
Requires-Dist: mkdocs-jupyter; extra == 'docs'
Requires-Dist: mkdocs-material; extra == 'docs'
Requires-Dist: mkdocstrings[python]; extra == 'docs'
Provides-Extra: ex
Requires-Dist: diffrax; extra == 'ex'
Requires-Dist: optax; extra == 'ex'
Description-Content-Type: text/markdown

# neojax

**neojax** (**Ne**ural **O**perators in JAX) is an implementation of Neural Operators built on top of [JAX](https://github.com/jax-ml/jax) and [Equinox](https://github.com/patrick-kidger/equinox). It provides a clean, modular API inspired by the original [neuraloperator](https://github.com/neuraloperator/neuraloperator) library.

Currently, **neojax** is in its early stages. Only the **Fourier Neural Operator (FNO)** is available, with support for:

- 1D, 2D, 3D, 4D, and 5D coordinate-based operator learning.
- Grid-based positional embeddings (`GridEmbeddingNd`).
- Symmetrical domain padding (`DomainPadding`).
- Pointwise MLP (Channel-MLP) expansions for improved expressivity.
- Various skip connections (Linear, Soft-Gating, Identity).
- Seamless integration with JAX's `vmap`, `jit`, and `grad`.

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

#### Installation
Install the python package via pypi
```bash
pip3 install neojax-operators
```

#### Quickstart
neojax exposes a similar API to neuraloperators and equinox and should therefore be familiar to use:
```python
import jax.numpy as jnp
from neojax.models import FNO

fno = FNO(n_modes=(64, 64),
        hidden_channels=64,
        in_channels=2,
        out_channels=1
    )

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

pred = fno(x)
```

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

#### Benchmarks
Coming soon (Please refer to `BENCHMARKS.md` for a performance comparison of `neojax` and `neuraloperator`.)

#### Motivation
Jax has become ubiquitous in Scientific Machine Learning (SciML) and Scientific Computing. This is largely due to its core design, which embraces mathematical and functional transformations (like `jit`, `vmap`, and `grad`) and seamlessly integrates with NumPy-like paradigms. However, despite Neural Operators fundamentally shaping the SciML landscape and being frequently used for solving PDEs, a comprehensive, native Jax implementation has been notably missing. `neojax` was created to bridge this gap, bringing the performance, predictability, and ecosystem of Jax to the Neural Operator community.

#### Design Choices
Although originally conceived as a direct port of the PyTorch `neuraloperator` library, `neojax` evolved into a ground-up, Jax-native re-implementation. This approach avoids the pitfalls of forcing PyTorch idioms into a functional framework and significantly reduces internal complexity. By building directly on `equinox`, `neojax` aligns perfectly with Jax's pure-functional design principles while maintaining a clean, accessible, and class-based API.

#### Roadplan
The first and currently only supported Neural Operator is a simple Fourier Neural Operator (FNO). In upcoming releases more models and components will be added.

#### 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.

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

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