foundax 0.1.0

One-stop JAX foundation model repository

foundax

One-stop JAX foundation model hub for this workspace.

The source repositories are vendored locally under repos/, so foundax can run against internal copies instead of the separate sibling repositories.

This repository provides one convert script and one compare script that dispatch to vendored model repositories under repos/:

• jax_morph
• jax_mpp
• jax_pdeformer2
• jax_poseidon
• jax_prose
• jax_walrus
• jax_dpot
• jax_bcat

Python API

You can import foundax as a package and instantiate models via:

import foundax

# MORPH-Ti
model = foundax.morph.Ti()

# MPP-Base
model = foundax.mpp.B(n_states=12)

# Poseidon-L
model = foundax.poseidon.L(num_in_channels=1, num_out_channels=1)

Version/size namespaces (mirroring jno/architectures/models.py):

• foundax.poseidon: T, B, L
• foundax.walrus: base (aliases: default, v1)
• foundax.morph: Ti, S, M, L
• foundax.mpp: Ti, S, B, L
• foundax.bcat: base (aliases: default, v1)
• foundax.pdeformer2: small, base, fast
• foundax.dpot: Ti, S, M, L, H
• foundax.prose: fd_1to1, fd_2to1, ode_2to1, pde_2to1

Lowercase aliases are also available for size variants where useful (for example foundax.morph.ti(), foundax.poseidon.b(), foundax.dpot.h()).

Quick Start

From foundax/:

python scripts/convert.py morph -- --input /path/to/model.pth --output /path/to/model.msgpack --model-size Ti
python scripts/compare.py morph -- --model-size Ti --checkpoint /path/to/model.pth

Vendored repositories live at:

foundax/repos/jax_morph
foundax/repos/jax_mpp
foundax/repos/jax_pdeformer2
foundax/repos/jax_poseidon
foundax/repos/jax_prose
foundax/repos/jax_walrus
foundax/repos/jax_bcat
foundax/repos/jax_dpot

General pattern:

python scripts/convert.py <model> [model-options...]
python scripts/compare.py <model> [model-options...]

If model options contain flags that overlap with foundax flags, put -- before model options.

Examples:

# DPOT
python scripts/convert.py dpot -- --input ../jax_dpot/ogrepo/model_Ti.pth --output ./artifacts/dpot_ti.msgpack
python scripts/compare.py dpot -- --checkpoint ../jax_dpot/ogrepo/model_Ti.pth

# PROSE ODE
python scripts/convert.py prose --prose-variant ode -- --input /path/prose_ode.pth --output /path/prose_ode.msgpack
python scripts/compare.py prose --prose-variant ode -- --checkpoint /path/prose_ode.pth --msgpack /path/prose_ode.msgpack

# BCAT (conversion supported directly in foundax)
python scripts/convert.py bcat -- --checkpoint /path/bcat.pt --output /path/bcat.msgpack

Notes

• compare for bcat is not available yet because there is no existing comparison script in jax_bcat.
• This repo intentionally reuses each vendored model repo's conversion/validation logic instead of duplicating that code.