JAXPIPhysics-informed neural networks at scale

A lean JAX library for physics-informed neural networks.

⚡

Multi-GPU by default

Data-parallel training via jax.shard_map with exact gradient averaging — the same script runs on 1 or N GPUs, bit-for-bit consistent.

🌀

Training algorithms that work

Grad-norm loss balancing, causal training, pseudo-time stepping, time-window curricula, and multi-stage homotopy — all toggled from the config.

🏴‍☠️

Modern architectures

PirateNets with residual adaptive blocks, modified MLPs, random Fourier features, and exact periodic embeddings.

🧪

16 benchmark examples

From 1D advection to 3D Taylor–Green and Kolmogorov flow at Re 10⁶ — each a self-contained script with configs and an evaluation notebook.

🧰

One-line trainer

train(config, model, sampler) handles stepping, adaptive weights, logging, checkpointing, and resume. Time windows are one call more.

✅

Tested core

A CPU-runnable test suite covers the sharded training step, causal weighting, checkpointing, and the trainer — including multi-device equivalence.

At a glance

Define your PDE residual, pick a config, and train:

python

from jaxpi.models import ForwardIVP, create_model
from jaxpi.samplers import UniformSampler
from jaxpi.training import train

class Burgers(ForwardIVP):
    def r_net(self, params, t, x):
        u = self.neural_net(params, t, x)
        u_t = grad(self.neural_net, argnums=1)(params, t, x)
        u_x = grad(self.neural_net, argnums=2)(params, t, x)
        u_xx = grad(grad(self.neural_net, argnums=2), argnums=2)(params, t, x)
        return u_t + u * u_x - 0.01 / jnp.pi * u_xx

model = create_model(config, Burgers, u0=u0, t_star=t_star, x_star=x_star)
train(config, model, UniformSampler(dom, config.training.batch_size))