Available Models
This page summarizes the foundation-model families currently exposed in jNO.
| Model family | jNO entry points | Expected input layout | Notes |
|---|---|---|---|
| MPP | jnn.nn.mppTi(...), jnn.nn.mppS(...), jnn.nn.mppB(...), jnn.nn.mppL(...) |
See synced upstream README | External model page synced from jax_mpp |
| Morph | jnn.nn.morphTi(...), jnn.nn.morphS(...), jnn.nn.morphM(...), jnn.nn.morphL(...) |
Structured image-like inputs; default spatial size is 128 and must be divisible by 8 |
MORPH family for arbitrary data modality |
| PDEformer2 | jnn.nn.pdeformer2_small(...), jnn.nn.pdeformer2_base(...), jnn.nn.pdeformer2_fast(...) |
See synced upstream README | External model page synced from jax_pdeformer2 |
| Poseidon | jnn.nn.poseidonT(...), jnn.nn.poseidonB(...), jnn.nn.poseidonL(...) |
Structured 2-D 128 x 128 grids, typically created with jno.domain.poseidon(nx=128, ny=128) |
Pretrained ScOT variants for PDE workloads |
| Walrus | jnn.nn.walrus(input_shape=..., num_out_channels=...) |
Channels-last tensors of shape (B, T, H, [W, [D]], C); commonly used with structured 128 x 128 grids |
Large foundation model using a SpaceBag encoder |
Grid Setup
For structured 2-D workflows, use the built-in Poseidon-style grid constructor:
This creates a structured grid matching the pixel layout expected by the foundation-model-oriented workflows in jNO. Physical groups are the same as for equi_distant_rect.