Loading...

jNO: JAX library for neural operator and PDE foundation model training

March 26, 2026

Authors

Leon ArmbrusterLeon Armbruster

jNO: JAX library for neural operator and PDE foundation model training


Scientific machine learning has reached a point where model quality is not the only bottleneck. In many projects, a substantial difficulty is workflow fragmentation: geometry handling, residual construction, supervised losses, optimization, and experiment management are often spread across separate abstractions. That fragmentation becomes more pronounced when researchers want to move between purely data-driven operator learning and physics-informed training.

jNO was developed with that problem in mind. At its core, it is mainly a symbolic tracing layer built above JAX for neural-operator and PDE-oriented workflows. It unifies domains, model calls, residual expressions, and training objectives within a single traced program. Rather than treating geometry, model definition, and PDE constraints as separate stages, jNO expresses them in one symbolic layer and compiles them into one optimization pipeline.

Motivation

Neural operators are increasingly used for surrogate modeling, inverse problems, and PDE-constrained learning. However, many practical workflows still require substantial orchestration code around the model itself. A researcher may need one API for mesh generation, another for loss terms, another for supervised learning, and yet another for tuning or deployment.

One practical goal of jNO is to reduce this separation. The library is designed so that users can define variables on domains, compose model outputs with differential operators, and train against both data and residual constraints without restructuring the surrounding code. In that sense, jNO is not simply a model zoo; it is an attempt to provide a more coherent programming model for scientific ML in JAX.

Core idea

At the center of jNO is a tracing-first design. Symbolic placeholders represent variables, constants, model calls, derivatives, and reductions. These objects can be composed algebraically to form residuals, supervised losses, and diagnostics. In practice, the library acts as a symbolic layer above JAX, with different JAX-native model implementations fitting into the same traced program. The resulting program is then compiled through JAX and XLA as a unified training computation.

This design has two practical consequences.

  • Data-driven and physics-informed objectives can be combined in the same workflow.
  • Mesh-aware operations, parameter conditioning, and model composition can be expressed without introducing a second orchestration layer.

For researchers working on PDE learning, this matters because project requirements often evolve. A workflow may begin with supervised operator regression, then later incorporate residual regularization, boundary-condition enforcement, or fine-tuning of a pretrained backbone. jNO is explicitly structured to support that progression.

Getting started

For installation, the package can be added with:

pip install jax-neural-operators

Example use case

One representative use case is parametric heat or diffusion modeling. Suppose a researcher wants to learn a mapping from spatial coordinates and a material parameter k to a field solution u(k, x, y). In a conventional setup, geometry sampling, neural network definition, PDE residual construction, and training may be handled by loosely connected components. In jNO, these pieces can be written as one traced program.

That is useful when the task requires both interpolation from data and consistency with governing equations. For example, a model can be trained on sampled solutions while also enforcing a diffusion-type residual over the same domain. The result is a workflow that remains compact while still accommodating mesh-aware evaluation and parameterized inputs.

Example workflow

The following example reflects a typical end-to-end usage pattern in jNO. To make the flow clearer, it is useful to separate it into the major stages of an operator-learning experiment.

1. Initialize the run

The workflow begins by importing the required libraries and creating a run directory. This provides a consistent location for outputs such as training curves and serialized model states.

import jno
import jax
import optax

dir = jno.setup("./runs/test")

2. Construct the domain and attach parameters

The next step defines a batched rectangular domain. Interior and boundary coordinates are extracted from the domain object, and a sample-wise parameter tensor k is attached. In this example, k can be interpreted as a spatially constant but sample-dependent coefficient for a diffusion-type equation.

dom = 500 * jno.domain.rect(mesh_size=0.05, x_range=(0, 2), y_range=(0, 1))
x, y, _ = dom.variable("interior")

k = dom.variable("k", jax.random.uniform(jax.random.PRNGKey(0), shape=(500, 1, 1), minval=0.5, maxval=1.5))

This stage illustrates one of the central ideas in jNO: the domain is not just a geometry container, but the object that organizes coordinates, tags, and parameterized inputs for the traced program.

3. Define the neural operator model

Once the domain variables are available, a DeepONet can be instantiated and equipped with an optimizer. Here, the network consumes the parameter tensor together with spatial coordinates.

net = jno.nn.deeponet(
	n_sensors=1,
	coord_dim=2,
	basis_functions=32,
	hidden_dim=128,
	activation=jax.numpy.tanh,
)
net.optimizer(
	optax.adam(
		learning_rate=optax.schedules.cosine_decay_schedule(
			init_value=1e-3, decay_steps=10_000, alpha=1e-5
		)
	)
)

The key point is that the model remains part of the same symbolic workflow as the domain and the residual expressions that follow.

4. Build the forward pass and PDE residual

The model output is then transformed so that it satisfies homogeneous boundary conditions through a multiplicative envelope. The PDE residual is defined immediately afterward using traced differential operators.

u = net(k, jno.np.concat([x, y], axis=-1)) * x * (2 - x) * y * (1 - y)
pde = k * (u.dd(x) + u.dd(y)) + 1.0

This is an important step conceptually. The forward pass, boundary-condition encoding, and residual definition all live in the same expression language. That can reduce the amount of glue code required to move from mathematics to training.

5. Create the solver, train, and save the result

After the residual has been specified, the solver is created from the traced constraints and the domain context. Training is then launched directly, and the trained state is serialized for later reuse.

crux = jno.core(constraints=[pde.mse], domain=dom, mesh=(1, 1)).print_shapes()
crux.solve(epochs=1_000, batchsize=32).plot(f"{dir}/training.png")
jno.save(crux, f"{dir}/model.pkl")

In practice, this stage brings together several tasks that are often separated in other workflows: loss assembly, runtime construction, optimization, logging, and persistence.

6. Evaluate on a finer test domain

Finally, the trained solver is evaluated on a new domain with a finer mesh and a different parameter range. This makes it possible to inspect the learned operator on unseen instances without redefining the full program.

tst_dom = 16 * jno.domain.rect(mesh_size=0.01)
tst_dom.variable(
	"k",
	jax.random.uniform(jax.random.PRNGKey(0), shape=(16, 1, 1), minval=0.1, maxval=1.9),
)

pred, x, y, k = crux.eval([u, x, y, k], domain=tst_dom)
print(pred.shape, x.shape, y.shape, k.shape)

Taken together, these steps show the intended appeal of jNO for scientific users. Domain definition, parameter handling, neural-operator construction, PDE residuals, training, persistence, and inference can all be expressed within one coherent workflow rather than being split across separate tool layers.

Architectural scope

jNO is broader than a single architecture. The library supports multiple neural-operator families, including spectral, convolutional, attention-based, branch-trunk, and point-based models. It also supports JAX-native adaptation workflows for PDE foundation-model families that have often been developed in PyTorch-centered ecosystems.

Specifically: jax_poseidon, jax_walrus, jax_morph, jax_pdeformer2 and jax_mpp

Equally important is the domain layer. jNO integrates mesh generation and loading, tagged regions such as interior and boundary subsets, tensor-valued conditioning variables, and finite-difference-aware metadata. This makes the framework suitable for workflows in which PDE residuals, spatial tags, and batched operator data all need to coexist.

The runtime and experiment layer extend that design further. jNO includes support for solver persistence, parameter-level optimizer control, masking and LoRA-style adaptation, hyperparameter tuning, and reusable serialized artifacts. From a research-software perspective, this is important because reproducibility and transferability are often as important as model accuracy.

Closing remarks

jNO represents a pragmatic step toward more unified scientific ML tooling in JAX. Its main strength is not a single model family, but the fact that it offers one programming interface for domains, residuals, supervision, and training. For researchers working on operator learning, hybrid physics-informed objectives, or PDE foundation-model transfer, that unification may reduce the amount of surrounding infrastructure required to run experiments.

As the scientific ML ecosystem matures, libraries that simplify these end-to-end workflows will likely remain useful. jNO is one example of that direction.

← Back to Blog