foundax 0.1.3

One-stop JAX foundation model repository

foundax

A unified JAX model zoo for neural operators, PINNs, and foundation models.

pip install foundax

Usage

import foundax as fx

# Neural operator architectures (returns initialized FlaxModel)
model = fx.mlp(in_features=2, output_dim=1, hidden_dims=64, num_layers=3)
model = fx.fno2d(in_features=1, hidden_channels=32, n_modes=16)
model = fx.unet2d(in_channels=1, out_channels=1)
model = fx.transformer(num_tokens=1000, d_model=128, num_heads=8)

# Poseidon — (B, 128, 128, C) → (B, 128, 128, C)
model = fx.poseidon.T()   # T/B/L variants
model = fx.poseidon.B()
model = fx.poseidon.L()

# MORPH — (B, t, F, C, D, H, W) → (B, F, C, D, H, W)
model = fx.morph.Ti()    # Ti/S/M/L variants
model = fx.morph.S()
model = fx.morph.M()
model = fx.morph.L()

# MPP — (T, B, C, H, W) → (B, C, H, W)
model = fx.mpp.Ti(n_states=12)  # Ti/S/B/L variants
model = fx.mpp.S(n_states=12)
model = fx.mpp.B(n_states=12)
model = fx.mpp.L(n_states=12)

# Walrus — (B, T, H, W, C) → (B, T, H, W, C_out)
model = fx.walrus.base()

# BCAT — (B, T_in+T_out, 128, 128, C) → (B, T_out, 128, 128, C)
model = fx.bcat.base()

# PDEformer-2 — graph inputs → (n_graph, n_points, 1)
model = fx.pdeformer2.small()  # small/base/fast variants
model = fx.pdeformer2.base()
model = fx.pdeformer2.fast()

# DPOT — (B, 128, 128, T, C) → (B, 128, 128, T_out, C)
model = fx.dpot.Ti()    # Ti/S/M/L/H variants
model = fx.dpot.S()
model = fx.dpot.M()
model = fx.dpot.L()
model = fx.dpot.H()

# PROSE — various sequence-to-sequence configurations
model, variables = fx.prose.fd_1to1()                          # (B, T_in, 128, 128, C) → (B, T_out, 128, 128, C)
model, variables = fx.prose.fd_2to1(n_words=64)                # (B, T_in, 128, 128, C) + symbols → predictions
model, variables = fx.prose.ode_2to1(n_words=64, pad_index=0)  # (T, 1, C) + text → (T_out, C)
model, variables = fx.prose.pde_2to1(n_words=64, pad_index=0)  # (T, 1, C) + text → spatial field

Integration with jNO

import foundax as fx
import jno

net = jno.nn.wrap(fx.mlp(in_features=2, output_dim=1))
net.optimizer(optax.adam, lr=1e-3)

License

This project is licensed under the MIT License.

Foundation models are subject to their original licenses. See THIRD_PARTY_LICENSES for details. Note that some pretrained weights (e.g. Poseidon) are released under non-commercial licenses.