Upload folder using huggingface_hub

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 MACE-MP
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

NOTICE

@@ -0,0 +1,25 @@
+MACE-MPA-0 (medium) — JAX re-export by CuspAI
+===============================================
+
+This artifact is a derivative work of MACE-MPA-0 medium.
+
+Original work
+-------------
+Copyright (c) 2024 MACE-MP
+Licensed under the MIT License. See LICENSE for the full text.
+Upstream: https://github.com/ACEsuit/mace-foundations
+
+Modifications
+-------------
+CuspAI re-exported the upstream PyTorch checkpoint
+(mace-mpa-0-medium.model) to a JAX computation graph using tojax
+(https://github.com/cusp-ai-oss/tojax). The model weights and architecture
+are unchanged, bit-for-bit, from the upstream release. Numerical agreement
+with the PyTorch reference was verified within rtol=1e-4, atol=1e-4 during
+export. The file format changed from .model to a .zip containing a
+jax.export graph plus msgpack-encoded parameters.
+
+Attribution
+-----------
+Please cite Batatia et al. (2024), J. Chem. Phys. 163, 184110
+(arXiv:2401.00096) when using this checkpoint in research.

README.md

@@ -0,0 +1,138 @@
+---
+license: mit
+license_link: LICENSE
+library_name: kups
+tags:
+  - chemistry
+  - materials-science
+  - molecular-dynamics
+  - interatomic-potential
+  - mlff
+  - jax
+  - mace
+---
+
+# MACE-MPA-0 medium — JAX build
+
+This repository hosts a JAX export of [MACE-MPA-0 medium](https://github.com/ACEsuit/mace-foundations) 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 the MACE team. 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 |
+|------|----------|---------|-------|
+| `mace-mpa-0-medium_32.zip` (40 MB) | [MACE-MPA-0 medium](https://github.com/ACEsuit/mace-foundations) | [MIT](LICENSE) | Batatia et al., *J. Chem. Phys.* **163**, 184110 (2024); [arXiv:2401.00096](https://arxiv.org/abs/2401.00096) |
+
+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-mace-jax",
+    filename="mace-mpa-0-medium_32.zip",
+)
+model = TojaxedMliap.from_zip_file(path)
+# model.cutoff, model.params, model.call(atom_graph_input) are now available.
+```
+
+kUPS ships CLI wrappers that take a YAML config pointing at this zip:
+
+```sh
+kups_relax_mlff --config relax_mace.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:** [ACEsuit/mace-foundations](https://github.com/ACEsuit/mace-foundations) · **Checkpoint:** `mace-mpa-0-medium.model`
+
+MACE-MPA-0 is a foundation-model member of the MACE family (Multi-ACE, a higher-body-order message-passing equivariant neural network), trained on the MPtrj + sAlex materials dataset covering 89 chemical elements. The "medium" variant is the default recommendation from the MACE team for general-purpose materials simulation.
+
+**Original authors:** Ilyes Batatia, Philipp Benner, Yuan Chiang, Alin M. Elena, Dávid P. Kovács, Janosh Riebesell, and the rest of the MACE team (~100 collaborators). Corresponding author: Gábor Csányi (Cambridge).
+
+**Supported elements:** atomic numbers 1–83, 89–94 (see `metadata.json` inside the zip).
+
+**Intended use and limitations:** general-purpose materials modelling (energies, forces, stresses) at DFT/PBE+U accuracy. Not trained for isolated molecules; for molecular systems refer to MACE-OFF. See the [upstream README](https://github.com/ACEsuit/mace-foundations) for authoritative guidance.
+
+**Citation:**
+
+```bibtex
+@article{batatia2024foundation,
+  title   = {A foundation model for atomistic materials chemistry},
+  author  = {Batatia, Ilyes and Benner, Philipp and Chiang, Yuan and Elena, Alin M.
+             and Kov{\'a}cs, D{\'a}vid P. and Riebesell, Janosh and others},
+  journal = {The Journal of Chemical Physics},
+  volume  = {163},
+  number  = {18},
+  pages   = {184110},
+  year    = {2024},
+  doi     = {10.1063/5.0230281},
+  eprint  = {2401.00096},
+  archivePrefix = {arXiv},
+}
+```
+
+## 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_mace.py --output mace-mpa-0-medium_32.zip --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 `.model` → 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 the MACE team. CuspAI's only contribution is the JAX export — we trained nothing, changed no weights, and designed none of the architecture. Please cite Batatia et al. (2024) when using this checkpoint in research.
+
+The file `mace-mpa-0-medium_32.zip` is distributed under the MIT License; see [LICENSE](LICENSE) for the full text and [NOTICE](NOTICE) for attribution and modification details. 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 MACE: please direct to the upstream authors via [ACEsuit/mace-foundations](https://github.com/ACEsuit/mace-foundations).

mace-mpa-0-medium_32.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5eb645f2dc2c904cfed1d364475202fdbb6309a5720de91e042dcb7d6ba7bf5
+size 39681963