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

jNO is a unified, JAX-native library that streamlines the training of neural operators and foundation models by integrating data-driven and physics-informed approaches into a single symbolic tracing system, enabling seamless transitions between operator regression, mesh-aware residual evaluation, and PDE-constrained optimization without code restructuring.

The Gist

Imagine you are trying to teach a computer to understand the laws of physics, like how heat flows through a metal plate or how water swirls around a rock. In the past, doing this with artificial intelligence was like trying to build a house where the architect, the plumber, the electrician, and the carpenter all spoke different languages and used different blueprints. You had to write one set of code for the shape of the room (geometry), another for the math equations (physics), and a third for the actual learning process. If you wanted to switch from one type of math to another, you often had to tear down the whole house and start over.

jNO (jax Neural Operators) is a new tool that acts like a "universal translator" and a "master builder" rolled into one. It is a software library designed to make training these physics-smart AI models much easier, faster, and more flexible, specifically for the JAX programming language (a popular tool for high-speed scientific computing).

Here is how it works, using some simple analogies:

1. The "One-Script" Magic (Tracing System)

Think of jNO as a single, magical script that controls the entire construction site.

• Before: You had to write a script for the blueprint, a separate script for the math, and another for the learning rules. If you wanted to change the math, you had to rewrite the blueprint script too.
• With jNO: You write everything in one language. You define the shape of the room, the physics equations, and the learning goals all in one go. The software "traces" (or records) your instructions like a movie director filming a scene. Later, it compiles this movie into a super-efficient, high-speed program. This means you can swap between different types of math problems or add new physics rules without rewriting your code.

2. The "Lego" Foundation Models

Currently, there are many different "foundation models" (pre-trained AI brains) for physics, but they are like Lego sets from different manufacturers that don't fit together. One brand uses red bricks, another uses blue, and they can't be stacked.

• jNO's Role: It acts as a universal adapter. It takes these different AI models (like Poseidon, Walrus, and Morph) and translates them so they all fit into the same JAX ecosystem. Now, a researcher can take a pre-trained "brain," tweak it slightly, and combine it with their own custom physics rules, all without needing to switch software tools.

3. The "Smart Mesh" (Handling Shapes)

When simulating physics, computers need to break shapes (like a curved pipe or a complex building) into tiny grid pieces called a "mesh."

• The Innovation: jNO has a built-in "smart mesh" system. It's like having a robot that can instantly draw a grid over any shape you describe, whether it's a simple square or a complex 3D object with holes. It keeps track of which part of the grid is the "inside," which is the "wall," and which is the "boundary," so the AI knows exactly where to apply the physics rules.

4. The "Fine-Tuning" Dial

Sometimes you want to take a pre-trained AI and teach it a specific new task.

• The Control Panel: jNO gives you a very detailed control panel. You can tell the AI, "Freeze these parts of your brain so they don't change," or "Only learn from these specific connections," or "Use a specific learning speed." You can do this for individual parts of the model without having to rebuild the whole thing. It's like being able to adjust the volume on just the drums in a song without changing the guitar or vocals.

5. The "Double-Mode" Engine (FEM and PINNs)

The paper highlights that jNO can handle two different ways of solving physics problems:

• Point-by-Point: Checking the physics at specific dots (like checking the temperature at specific spots on a map).
• Whole-Shape (Finite Element): Looking at the physics as a continuous flow over a whole shape (like calculating the total stress on a bridge).
• The Benefit: jNO lets you switch between these two modes using the same code. It's like having a car that can drive on both dirt roads and highways without you needing to change the engine or the steering wheel.

Why Does This Matter?

The main goal of jNO is to stop the "fragmentation" of scientific software. Instead of researchers juggling five different tools to train one AI model, jNO brings everything into one place.

• Speed: Because it uses JAX's special compilation features, it runs faster on modern computer chips.
• Simplicity: You don't need to be a software architect to switch between different types of physics problems.
• Reusability: Once you write a program in jNO, you can save it, share it, and run it again later, even on different computers, with confidence that it will work the same way.

In short, jNO is trying to make the complex world of "scientific machine learning" feel as simple and unified as writing a single, coherent story, rather than stitching together a patchwork quilt of different code fragments.

Technical Summary

Technical Summary: jNO – A JAX Library for Neural Operator and Foundation Model Training

Problem Statement

The current landscape of scientific machine learning (SciML), particularly regarding neural operators and Physics-Informed Neural Networks (PINNs), suffers from significant software fragmentation. State-of-the-art foundation models (e.g., Poseidon, Walrus, PDEformer2) and operator learning tools are predominantly maintained in separate repositories within the PyTorch ecosystem. This separation creates barriers for systematic comparison, transfer learning, and reproducible fine-tuning across different model families. Furthermore, practical workflows are often disjointed, requiring users to manage separate abstractions for geometry, loss construction, model definition, and training logic. This fragmentation is exacerbated when attempting to combine supervised operator learning with physics-informed residuals or when transferring pretrained backbones into new Partial Differential Equation (PDE) settings.

Methodology

jNO (jax Neural Operators) addresses these challenges by providing a unified, JAX-native library that treats the entire training workflow as a single traced program. The core methodology relies on a symbolic Domain-Specific Language (DSL) centered on tracing and deferred execution.

1. Tracing-First Design

Unlike eager execution models, jNO constructs a deferred computation graph where domains, model calls, residuals, supervised losses, and diagnostics are written in one symbolic language.

• Symbolic Nodes: The system uses Placeholder nodes to represent variables, tensors, models, and differential operators. Arithmetic and differential operations overload standard JAX primitives to build expression trees rather than executing immediately.
• Unified Pipeline: Users define domains and variables, compose models with differential expressions, and solve everything through a single interface. This allows seamless transitions between operator regression, mesh-aware residual evaluation, and PDE-constrained training without restructuring code.
• Optimization: Before execution, the system performs common sub-expression elimination (CSE) on the traced tree, canonicalizing structurally identical subtrees to a single XLA program. This unifies objectives, operator definitions, and physics constraints.

2. Domain and Mesh Workflow

The domain class acts as a mesh-aware data organization layer, bridging geometry definitions and traced expressions.

• Integration: It supports mesh generation via PyGmsh/Gmsh and loading via meshio.
• Structure: Data is organized into a mesh pool (for tagged point sets) and a runtime context dictionary. It supports batched training where each batch element can represent a different PDE parameterization or forcing field.
• Adaptivity: The system includes hooks for adaptive point management, allowing training to update point sets without altering symbolic equations.

3. Finite Element Method (FEM) Integration

jNO extends its traced programming model to weak form assembly via the FEAX backend.

• Unified DSL: Instead of a separate FEM front-end, jNO uses the same symbolic layer for mesh-based weak forms, variational PINNs, and operator learning.
• Assembly: Symbolic weak forms are lowered to backend-neutral representations. Depending on the target, the system can dispatch to vpinn (variational PINN), fem_system (linear steady), fem_residual (nonlinear steady), or fem_time (transient).
• Differentiability: The integration keeps assembled finite element operators within the JAX ecosystem, enabling end-to-end differentiable workflows (assembly, residual evaluation, Jacobian construction) compatible with JAX transformations.

4. Model Interface and Foundation Models

jNO treats neural operator architectures as ordinary operations within the traced program.

• Foundation Models: It supports JAX-native translations of PDE foundation models (Poseidon, Walrus, Morph, etc.) maintained in the external foundax repository. These models can be wrapped via jno.nn.wrap() and composed with arbitrary traced expressions.
• Fine-Grained Control: Users can configure model, optimizer, and learning rate behaviors at the parameter level. Features include selective parameter masking, freezing/unfreezing, and LoRA (Low-Rank Adaptation) support for both Equinox and Flax models.

5. Execution and Training

The jno.core module compiles the traced constraints into a single JIT-compiled training program.

• Performance: It leverages XLA for aggressive sub-expression elimination, memory optimization, and multi-device parallelism (via a 2-axis device mesh).
• Hyperparameter Tuning: The ArchSpace interface supports grid search and gradient-free optimization (via Nevergrad) for architecture and training parameters, integrated directly into the training loop.
• Persistence: Experiments can be serialized using jno.save (supporting RSA signatures for integrity) to preserve solver states, domain objects, and compiled inference wrappers.

Key Contributions

1. Unified Traced Language: jNO is the first consolidated framework to express neural operators, physics-informed residuals, and foundation model fine-tuning within a single symbolic interface in JAX.
2. Ecosystem Consolidation: It acts as a unified backend for translating and integrating diverse PDE foundation model families (e.g., Poseidon, Walrus) into JAX, enabling side-by-side evaluation and composition without maintaining separate codebases.
3. FEM and Operator Unification: By extending the traced DSL to weak form assembly, jNO allows mesh-based weak forms, variational PINNs, and operator models to be composed without changing the programming model.
4. Parameter-Level Control: The library provides fine-grained control over training dynamics, including selective masking and LoRA adaptation, directly within the symbolic flow.
5. Reproducibility and Portability: The library is packaged for PyPI, includes automated tests, and supports signed artifact sharing for verifiable workflows.

Results and Claims

The paper does not present new benchmark results on specific PDE datasets but focuses on the architectural and workflow contributions:

• Workflow Efficiency: The paper claims that the traced design eliminates the need for code restructuring when switching between operator learning, residual evaluation, and PDE-constrained training.
• Performance: By compiling to a single XLA graph, jNO claims to reduce compilation overhead and improve numerical stability compared to PyTorch-centric alternatives, while maintaining prototyping simplicity.
• Interoperability: The integration of FEAX allows for differentiable finite element workflows that remain compatible with JAX solvers (e.g., Diffrax, Optimistix) or external libraries (e.g., SciPy).

Significance

The authors position jNO as a critical step toward reducing fragmentation in the SciML ecosystem. Its significance lies in:

• Lowering Barriers: It makes advanced neural operator and foundation model workflows more accessible to researchers by providing a unified, JAX-native stack.
• Enabling Hybrid Workflows: It facilitates previously difficult workflows, such as ensemble methods, progressive transfer learning, and multi-model constraints, by standardizing training controls and parameter-level optimization.
• Foundation for Future Research: By consolidating foundation models into a single framework, it enables systematic comparison and compositional workflows that were previously impeded by disparate repositories.

The paper concludes that jNO provides a credible research significance by making these advanced workflows more accessible and efficient within the JAX ecosystem, serving as a reusable research software stack for the community.