File size: 5,948 Bytes

07a254a

---
license: apache-2.0
library_name: kups
tags:
  - chemistry
  - materials-science
  - molecular-dynamics
  - interatomic-potential
  - mlff
  - jax
  - orb
---

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

This repository hosts a JAX export of [Orb v3 conservative-inf-omat](https://github.com/orbital-materials/orb-models) for use with [kUPS](https://github.com/cusp-ai-oss/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](https://github.com/cusp-ai-oss/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](https://github.com/orbital-materials/orb-models) | [Apache 2.0](LICENSE) | Rhodes et al. (2025), [arXiv:2504.06231](https://arxiv.org/abs/2504.06231) |

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

## Quick start

```sh
pip install kups[cuda]
```

```python
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:

```sh
kups_md_mlff --config md_orb.yaml
```

Example configs live in the kUPS [examples/](https://github.com/cusp-ai-oss/kups/tree/main/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](https://github.com/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](https://github.com/orbital-materials/orb-models#models) for authoritative guidance, including trade-offs between the `conservative`, `direct`, and `inf` variants.

**Citation:**

```bibtex
@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/](https://github.com/cusp-ai-oss/tojax):

```sh
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](LICENSE) for the full text and [NOTICE](NOTICE) for attribution and the §4(b) modification statement. The kUPS / tojax tooling citations:

```bibtex
@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](https://github.com/cusp-ai-oss/kups/issues)
- Scientific questions about Orb: please direct to the upstream authors via [orbital-materials/orb-models](https://github.com/orbital-materials/orb-models).