{"id":46695451,"url":"https://github.com/bamescience/mace-jax-parity","last_synced_at":"2026-03-09T05:04:12.680Z","repository":{"id":335076525,"uuid":"1097637347","full_name":"BAMeScience/mace-jax-parity","owner":"BAMeScience","description":null,"archived":false,"fork":false,"pushed_at":"2026-01-28T12:37:46.000Z","size":16803,"stargazers_count":3,"open_issues_count":0,"forks_count":0,"subscribers_count":0,"default_branch":"main","last_synced_at":"2026-01-29T01:10:05.396Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/BAMeScience.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2025-11-16T15:03:21.000Z","updated_at":"2026-01-28T12:38:17.000Z","dependencies_parsed_at":null,"dependency_job_id":null,"html_url":"https://github.com/BAMeScience/mace-jax-parity","commit_stats":null,"previous_names":["bamescience/mace-jax-parity"],"tags_count":null,"template":false,"template_full_name":null,"purl":"pkg:github/BAMeScience/mace-jax-parity","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BAMeScience%2Fmace-jax-parity","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BAMeScience%2Fmace-jax-parity/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BAMeScience%2Fmace-jax-parity/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BAMeScience%2Fmace-jax-parity/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/BAMeScience","download_url":"https://codeload.github.com/BAMeScience/mace-jax-parity/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/BAMeScience%2Fmace-jax-parity/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":30283703,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-03-09T02:57:19.223Z","status":"ssl_error","status_checked_at":"2026-03-09T02:56:26.373Z","response_time":61,"last_error":"SSL_read: unexpected eof while reading","robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":false,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2026-03-09T05:04:09.929Z","updated_at":"2026-03-09T05:04:12.666Z","avatar_url":"https://github.com/BAMeScience.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# mace-jax-parity\n\nUtilities to load or download a Torch MACE foundation model, convert it to a JAX bundle, compare their energies on mp-traj, and benchmark throughput.\n\n```\nscripts/\n├─ compare_mace_torch_jax.py      # parity check with progress bars and CSV logging\n├─ benchmark_mace_torch_predict.py # Torch inference benchmark (Accelerate, CSV output)\n├─ benchmark_mace_jax_predict.py   # JAX benchmark (compile time + throughput)\n├─ benchmark_mace_jax_train.py    # JAX training benchmark (one epoch w/ optax + CSV)\n├─ convert_mace_model_to_jax.py   # torch → JAX bundle converter (pass --dtype to control casting)\n├─ create_mace_foundation_model.py\n├─ check_cueq_torch.py            # report whether a Torch checkpoint uses cuEq kernels\n├─ check_cueq_jax.py              # inspect a JAX bundle for stored cuEq config\n├─ plot_energy_diff.py            # CPU/GPU relative ΔE histograms (log scale)\n```\n\nOther important directories:\n\n- `data/mptraj/` — contains `train.h5` / `valid.h5` subsets.\n- `models/` — generated artifacts (dtype-specific Torch checkpoints and MACE-JAX bundles), e.g.\n  - `mace_foundation_f32.pt`, `mace_foundation_f64.pt`\n  - `mace_jax_bundle_f32/`, `mace_jax_bundle_f64/`\n- `results/` — CSVs and plots emitted by the Makefile targets.\n- `makefile` — orchestration for model creation, parity checks (float32/float64), benchmarks, and plotting.\n\n## Quickstart\nAssumes the virtualenv at `/home/pbenner/Env/mace-jax/.venv` is activated.\n\n```bash\ncd mace-jax-parity\nsource ../.venv/bin/activate\n\n# Build Torch foundation models and JAX bundles (float32 + float64).\nmake models/mace_jax_bundle_f32 models/mace_jax_bundle_f64\n\n# Run comparisons (float32 + float64) and produce plots.\nmake compare\nmake plot-comparison\n\n# Run benchmarks and write CSV summaries.\nmake benchmark\n\n# Direct script usage examples:\npython scripts/compare_mace_torch_jax.py \\\n  --torch-model models/mace_foundation_f64.pt \\\n  --jax-model models/mace_jax_bundle_f64 \\\n  --data-dir data/mptraj \\\n  --split valid \\\n  --dtype float64 \\\n  --device cuda \\\n  --diff-csv results/custom_compare.csv \\\n  --tqdm\n\npython scripts/benchmark_mace_jax_predict.py \\\n  --torch-model models/mace_foundation_f32.pt \\\n  --jax-model models/mace_jax_bundle_f32 \\\n  --data-dir data/mptraj \\\n  --split valid \\\n  --dtype float32 \\\n  --device cuda \\\n  --multi-gpu \\\n  --max-edges-per-batch 480000 \\\n  --csv-output results/benchmark_jax.csv\n\npython scripts/benchmark_mace_jax_train.py \\\n  --jax-model models/mace_jax_bundle_f32 \\\n  --data-dir data/mptraj \\\n  --split train \\\n  --dtype float32 \\\n  --batch-size 16 \\\n  --learning-rate 1e-3 \\\n  --multi-gpu \\\n  --max-edges-per-batch 480000 \\\n  --csv-output results/benchmark_jax_train.csv\n```\n\n## Notes\n- `make compare` now runs both float32 and float64 suites (`results/compare_*_f32.csv` / `_f64.csv`). Plotting creates separate figures per dtype.\n- `compare_mace_torch_jax.py` forces both Torch and JAX onto the same device (`--device`) and shows tqdm for each backend. CSV output includes absolute and relative ΔE for further analysis.\n- Benchmarks write machine-readable rows (backend, dtype, device, graphs/s, wall time, compile time). Torch uses Accelerate for multi-GPU, JAX reports XLA compile overhead. The training benchmark shares the same CSV-friendly format and reports per-epoch losses.\n- `plot_energy_diff.py` consumes the CSVs and produces log-scale histograms of relative |ΔE|/scale for CPU vs GPU comparisons; pass `--dtype` to annotate the plots.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fbamescience%2Fmace-jax-parity","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fbamescience%2Fmace-jax-parity","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fbamescience%2Fmace-jax-parity/lists"}