# jNO: A JAX Library for Neural Operator and Foundation Model Training

arXiv: https://arxiv.org/abs/2605.10159 · HTML: https://mlsanity.com/paper/2605.10159/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 11 of its 26 references that are in the map.
- First posted on arXiv 11 May 2026.

## Builds on (most cited first, 11 of 11)

- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [DeepXDE: A deep learning library for solving differential equations](https://mlsanity.com/paper/1907.04502/) (2019)
- [Neural Operators for Accelerating Scientific Simulations and Design](https://mlsanity.com/paper/2309.15325/) (2023)
- [DPOT: Auto-Regressive Denoising Operator Transformer for Large-Scale PDE Pre-Training](https://mlsanity.com/paper/2403.03542/) (2024)
- [Variational Physics-Informed Neural Networks For Solving Partial Differential Equations](https://mlsanity.com/paper/1912.00873/) (2019)
- [Equinox: neural networks in JAX via callable PyTrees and filtered transformations](https://mlsanity.com/paper/2111.00254/) (2021)
- [BCAT: A Block Causal Transformer for PDE Foundation Models for Fluid Dynamics](https://mlsanity.com/paper/2501.18972/) (2025)
- [MORPH: PDE Foundation Models with Arbitrary Data Modality](https://mlsanity.com/paper/2509.21670/) (2025)
- [Physics-informed fine-tuning of foundation models for partial differential equations](https://mlsanity.com/paper/2603.15431/) (2026)

## Abstract (arXiv)

jNO (jax Neural Operators) is a JAX-native library for neural operators and foundation models with unified support for both data-driven and physics-informed training. Its core design is a tracing system in which domains, model calls, residuals, supervised losses, and diagnostics are written in one symbolic language and compiled into one optimization pipeline. This allows users to move between operator regression, mesh-aware residual evaluation, and PDE-constrained training without restructuring the surrounding code. jNO also supports multi-model compositions, fine-grained control at parameter level (model, optimizer, and learning rate), hyperparameter tuning, and JAX-native workflows for translated PDE foundation-model families. The source repository is available at https://github.com/FhG-IISB/jNO.