CuspAI/kUPS-orb-jax

Orb v3 (conservative-inf-omat) — JAX build

This repository hosts a JAX export of Orb v3 conservative-inf-omat for use with kUPS, a JAX-native molecular-simulation toolkit. The artefact is a self-contained .zip containing the serialized JAX computation graph, the original model parameters, and the minimal metadata needed to run inference.

Important: this is a re-export, not a retraining. Weights and architecture are the originals released by Orbital Materials. CuspAI's only contribution is converting the PyTorch reference implementation to JAX via tojax. Every scientific claim, citation, and credit belongs to the original authors.

Included model

File	Upstream	License	Paper
orb_v3_conservative_inf_omat.zip (102 MB)	Orb v3 conservative-inf-omat	Apache 2.0	Rhodes et al. (2025), arXiv:2504.06231

Cutoff radius 6.0 Å. Schema: AtomGraphInput (positions, atomic numbers, cell, pbc, edge index, cell offsets, batch index, charge, spin).

Quick start

pip install kups[cuda]

from huggingface_hub import hf_hub_download
from kups.potential.mliap.tojax import TojaxedMliap

path = hf_hub_download(
    repo_id="CuspAI/kUPS-orb-jax",
    filename="orb_v3_conservative_inf_omat.zip",
)
model = TojaxedMliap.from_zip_file(path)

kUPS ships CLI wrappers that take a YAML config pointing at this zip:

kups_md_mlff --config md_orb.yaml

Example configs live in the kUPS examples/ directory.

What's in the .zip

• model.jax — JAX computation graph, serialized via jax.export.
• params.msgpack — parameters as a msgpack-encoded list of arrays.
• metadata.json — cutoff radius and supported atomic numbers.
• dtypes.json — input dtypes for AtomGraphInput.

Exported with symbolic shapes (--symbolic NSE); accepts variable atom, system, and edge counts without recompilation.

Model details

Upstream: orbital-materials/orb-models · Checkpoint: orb_v3_conservative_inf_omat

Orb v3 is a non-equivariant graph network from Orbital Materials. The conservative_inf_omat variant computes forces and stresses by back-propagation through the energy (i.e. energy-conserving), uses unlimited neighbours per atom, and is trained on the OMat24 dataset (~55 M structures, modern VASP PBE54 pseudopotentials).

Original authors: Benjamin Rhodes, Sander Vandenhaute, Vaidotas Šimkus, James Gin, Jonathan Godwin, Tim Duignan, Mark Neumann. Orbital Materials, Inc.

Intended use and limitations: intended for inorganic materials at DFT (OMat24) accuracy. Confidence estimates are produced per-atom by the model. See the upstream model cards for authoritative guidance, including trade-offs between the conservative, direct, and inf variants.

Citation:

@misc{rhodes2025orbv3,
  title         = {Orb-v3: atomistic simulation at scale},
  author        = {Rhodes, Benjamin and Vandenhaute, Sander and {\v S}imkus, Vaidotas
                   and Gin, James and Godwin, Jonathan and Duignan, Tim and Neumann, Mark},
  year          = {2025},
  eprint        = {2504.06231},
  archivePrefix = {arXiv},
  primaryClass  = {cond-mat.mtrl-sci},
}

Export pipeline and reproducibility

The archive was produced with the exporter in tojax/examples/mlff/:

uv run python export_orb.py --output orb_v3_conservative_inf_omat.zip \
    --model-name orb_v3_conservative_inf_omat --symbolic NSE

tojax's export harness verifies numerical agreement with the PyTorch reference (default tolerances rtol=1e-4, atol=1e-4) before saving the archive.

Changes from upstream

• File format. PyTorch .pt → JAX-exported .zip (graph + msgpack params).
• Weights. Unchanged, bit-for-bit, from the upstream checkpoint.
• Architecture. Translated operation-for-operation; no approximations or substitutions.
• Numerics. Verified within rtol=1e-4, atol=1e-4 against the PyTorch reference during export.
• Cutoff and dtypes. Preserved from upstream defaults (6.0 Å; float32 positions/cell/offsets, int64 indices, bool PBC).
• Symbolic shapes. Exports accept variable numbers of atoms, systems, and edges without recompilation.

Attribution and license

This model exists because of the work of Orbital Materials. CuspAI's only contribution is the JAX export — we trained nothing, changed no weights, and designed none of the architecture. Please cite Rhodes et al. (2025) when using this checkpoint in research.

The file orb_v3_conservative_inf_omat.zip is distributed under the Apache License, Version 2.0; see LICENSE for the full text and NOTICE for attribution and the §4(b) modification statement. The kUPS / tojax tooling citations:

@software{kups2026,
  author = {{CuspAI}},
  title  = {kUPS},
  year   = {2026},
  url    = {https://github.com/cusp-ai-oss/kups},
}

@software{tojax2026,
  author = {{CuspAI}},
  title  = {tojax},
  year   = {2026},
  url    = {https://github.com/cusp-ai-oss/tojax},
}

Contact

• Issues with the JAX export or with kUPS: github.com/cusp-ai-oss/kups/issues
• Scientific questions about Orb: please direct to the upstream authors via orbital-materials/orb-models.