lammps-jax 0.2.0__tar.gz

lammps_jax-0.2.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Abhijeet Gangan
+
+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.

lammps_jax-0.2.0/NOTICE

@@ -0,0 +1,13 @@
+lammps-jax is licensed under the MIT license (see LICENSE).
+
+The following files are vendored from OpenXLA
+(https://github.com/openxla/xla), Copyright 2022-2024 The OpenXLA
+Authors, licensed under the Apache License, Version 2.0
+(http://www.apache.org/licenses/LICENSE-2.0):
+
+  cpp/third_party/pjrt_c_api.h            (xla/pjrt/c/pjrt_c_api.h)
+  cpp/third_party/pjrt_c_api_ffi_extension.h
+                                          (xla/pjrt/c/pjrt_c_api_ffi_extension.h)
+
+The only modification is the pjrt_c_api.h include path in the extension
+header; the C API header is unmodified apart from a provenance note.

lammps_jax-0.2.0/PKG-INFO

@@ -0,0 +1,187 @@
+Metadata-Version: 2.4
+Name: lammps-jax
+Version: 0.2.0
+Summary: CUDA-first LAMMPS/Kokkos plugin support for exported JAX models
+License-Expression: MIT
+Project-URL: Repository, https://github.com/abhijeetgangan/lammps-jax
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+License-File: NOTICE
+Requires-Dist: jax>=0.8
+Requires-Dist: jaxlib!=0.10.2,>=0.8
+Provides-Extra: cpu
+Requires-Dist: lammps-jax-pjrt-cpu>=0.11.2; (sys_platform == "linux" and platform_machine == "x86_64") and extra == "cpu"
+Provides-Extra: dist
+Requires-Dist: ase>=3.22; extra == "dist"
+Requires-Dist: e3nn-jax>=0.20; extra == "dist"
+Requires-Dist: jax!=0.11.2; extra == "dist"
+Requires-Dist: jax-md>=0.2.27; extra == "dist"
+Requires-Dist: nequix>=0.4; extra == "dist"
+Provides-Extra: test
+Requires-Dist: pytest>=7; extra == "test"
+Requires-Dist: numpy>=1.24; extra == "test"
+Dynamic: license-file
+
+# LAMMPS-JAX
+
+[![DOI](https://zenodo.org/badge/1278620240.svg)](https://doi.org/10.5281/zenodo.20838693)
+
+[LAMMPS](https://www.lammps.org) plugin for running JAX forcefields with
+`jax.export` and a PJRT runtime.
+
+## Setup
+
+The Python package is on PyPI; the LAMMPS plugin and the examples build from
+this repository.
+
+```bash
+pip install lammps-jax 'jax[cuda12]'  # GPU
+pip install 'lammps-jax[cpu]'         # CPU, with the XLA CPU PJRT plugin
+```
+
+From a checkout:
+
+```bash
+git clone https://github.com/abhijeetgangan/lammps-jax && cd lammps-jax
+uv venv .venv --python python3
+source .venv/bin/activate
+uv pip install -e '.[test]'
+uv pip install -U 'jax[cuda12]'
+```
+
+## Build
+
+The pair style needs the KOKKOS precision layer from LAMMPS 10 Sep 2025 or
+newer; stable releases do not carry it yet, and older trees stop at a
+compile-time error. Build the plugin against the same LAMMPS source tree as
+the `lmp` binary that loads it, and rebuild both together after pulling.
+
+```bash
+cmake -S cpp -B build-plugin-gpu-pjrt \
+  -D CMAKE_CXX_COMPILER="$KOKKOS_NVCC_WRAPPER" \
+  -D CMAKE_BUILD_TYPE=Release \
+  -D CMAKE_CXX_FLAGS="-fno-lto" \
+  -D CMAKE_SHARED_LINKER_FLAGS="-fno-lto" \
+  -D CMAKE_INTERPROCEDURAL_OPTIMIZATION=OFF \
+  -D LAMMPS_HEADER_DIR="$LAMMPS_SRC/src" \
+  -D JAXLIB_INCLUDE_DIR="$JAXLIB_INCLUDE_DIR" \
+  -D KOKKOS_CONFIG_INCLUDE_DIR="$LAMMPS_BUILD/lib/kokkos"
+cmake --build build-plugin-gpu-pjrt -j4
+```
+
+## Run
+
+```bash
+python examples/export_model.py lj examples/lj.lammps-jax.json \
+  --max-atoms 1024 --edges-per-atom 96
+```
+
+The bundle loads as a pair style:
+
+```
+pair_style jax/kk ${pjrt}
+pair_coeff * * examples/lj.lammps-jax.json
+```
+
+```bash
+LAMMPS_PLUGIN_PATH=build-plugin-gpu-pjrt \
+  "$LAMMPS_INSTALL/bin/lmp" -k on g 1 -pk kokkos newton off neigh full -sf kk \
+  -var pjrt /absolute/path/xla_cuda_plugin.so -in examples/in.lj_jax
+```
+
+## CPU
+
+Kokkos OpenMP builds run bundles exported with `platform="cpu"` on the XLA CPU
+PJRT plugin from the `cpu` extra (Linux x86-64):
+
+```bash
+uv pip install -e '.[cpu]'
+python -c 'import lammps_jax_pjrt_cpu as p; print(p.PLUGIN_PATH)'
+```
+
+Build LAMMPS with `-C cmake/presets/kokkos-openmp.cmake -D PKG_PLUGIN=on
+-D BUILD_MPI=on -D Kokkos_ENABLE_IMPL_VIEW_LEGACY=on` and the plugin as above
+against that build, without `CMAKE_CXX_COMPILER`. One MPI rank per core runs
+fastest:
+
+```bash
+python examples/export_model.py eam cuzr.lammps-jax.json \
+  --setfl examples/potentials/CuZr.eam.alloy.gz \
+  --max-atoms 8192 --edges-per-atom 128 --precision float64 --platform cpu
+LAMMPS_PLUGIN_PATH=build-plugin-cpu-pjrt \
+  mpirun -np 4 "$LAMMPS_CPU_INSTALL/bin/lmp" -k on t 1 -sf kk -pk kokkos newton on neigh half \
+  -var pjrt /absolute/path/pjrt_c_api_cpu_plugin.so -in examples/in.eam_cuzr
+```
+
+cuEquivariance and OpenEquivariance kernels are CUDA only.
+
+## Test
+
+```bash
+JAX_PLATFORMS=cpu python -m pytest
+```
+
+Integration tests, which require a GPU and LAMMPS:
+
+```bash
+LAMMPS_BIN=$LAMMPS_INSTALL/bin/lmp \
+PJRT_PLUGIN=/absolute/path/xla_cuda_plugin.so \
+LAMMPS_PLUGIN_PATH=$PWD/build-plugin-gpu-pjrt \
+  python -m pytest tests/test_lammps.py -v
+```
+
+For a host Kokkos LAMMPS and the CPU plugin, add `LAMMPS_JAX_TEST_PLATFORM=cpu`.
+
+## Design notes
+
+Models are exported with `jax.export` into a JSON bundle holding the
+serialized program and the settings the pair style enforces at load: atom and
+edge capacities, cutoff, precision, unit style, and the distribution scheme.
+Exported programs have static shapes, so positions and the edge list are
+padded to capacity and an edge mask marks the live entries. The pair style
+distance-filters pairs to the model cutoff and repacks the edge list every
+step, so edge capacity sizes to the cutoff rather than cutoff plus skin.
+
+The pair style executes the bundle through the PJRT C API, loading the same
+plugin library jax uses on the GPU, so the build never needs XLA. Kokkos kernels
+pack the LAMMPS neighbor list into sender, receiver, and mask arrays on
+device, the program runs on the LAMMPS CUDA stream, and forces come back as
+device buffers, so array data never passes through the host.
+
+Multi-rank runs pick one of two schemes at export. A ghost export widens the
+ghost shell to n_hops cutoffs and stays communication-free inside the
+program; a comm export keeps the one-cutoff shell and exchanges per-layer
+features through LAMMPS forward and reverse communication, called from inside
+the program as an FFI callback. Exchanges are linear JAX primitives whose
+transposes swap forward and reverse, so models may call either direction and
+every differentiation mode traces through them. Communicating half-edge
+exports pack each pair on one rank, matching the native half neighbor list,
+and reverse-communicate ghost density partials into their owners. Exchanged
+rows move in place on device under the Kokkos brick and tiled comm styles and
+stage through pinned host memory otherwise.
+
+cuEquivariance and OpenEquivariance kernels stay in the exported program as
+custom call targets, resolved at run time from LAMMPS_JAX_FFI_HANDLERS;
+contrib/ffi-replay handles libraries whose compiled kernels only exist inside
+the exporting process. NequIP checkpoints export through the uv-runnable
+scripts in examples/torch, which trace the torch blocks via torchax into
+the serialized program; the bundles load and run like any native export.
+
+## Citation
+
+If you use `lammps-jax` or any algorithms implemented here, please cite
+the archived Zenodo release: https://doi.org/10.5281/zenodo.20838694
+
+## Acknowledgements
+
+Thanks to @mitkotak for API discussions and @wcwitt for guidance on
+distributed inference. See
+[lammps/lammps#4691](https://github.com/lammps/lammps/pull/4691) for the
+original discussion on running JAX models in LAMMPS, and
+[openmm-jax](https://github.com/atomicarchitects/openmm-jax) for the OpenMM
+counterpart.
+
+## License
+
+MIT. See `LICENSE`.

lammps_jax-0.2.0/README.md

@@ -0,0 +1,162 @@
+# LAMMPS-JAX
+
+[![DOI](https://zenodo.org/badge/1278620240.svg)](https://doi.org/10.5281/zenodo.20838693)
+
+[LAMMPS](https://www.lammps.org) plugin for running JAX forcefields with
+`jax.export` and a PJRT runtime.
+
+## Setup
+
+The Python package is on PyPI; the LAMMPS plugin and the examples build from
+this repository.
+
+```bash
+pip install lammps-jax 'jax[cuda12]'  # GPU
+pip install 'lammps-jax[cpu]'         # CPU, with the XLA CPU PJRT plugin
+```
+
+From a checkout:
+
+```bash
+git clone https://github.com/abhijeetgangan/lammps-jax && cd lammps-jax
+uv venv .venv --python python3
+source .venv/bin/activate
+uv pip install -e '.[test]'
+uv pip install -U 'jax[cuda12]'
+```
+
+## Build
+
+The pair style needs the KOKKOS precision layer from LAMMPS 10 Sep 2025 or
+newer; stable releases do not carry it yet, and older trees stop at a
+compile-time error. Build the plugin against the same LAMMPS source tree as
+the `lmp` binary that loads it, and rebuild both together after pulling.
+
+```bash
+cmake -S cpp -B build-plugin-gpu-pjrt \
+  -D CMAKE_CXX_COMPILER="$KOKKOS_NVCC_WRAPPER" \
+  -D CMAKE_BUILD_TYPE=Release \
+  -D CMAKE_CXX_FLAGS="-fno-lto" \
+  -D CMAKE_SHARED_LINKER_FLAGS="-fno-lto" \
+  -D CMAKE_INTERPROCEDURAL_OPTIMIZATION=OFF \
+  -D LAMMPS_HEADER_DIR="$LAMMPS_SRC/src" \
+  -D JAXLIB_INCLUDE_DIR="$JAXLIB_INCLUDE_DIR" \
+  -D KOKKOS_CONFIG_INCLUDE_DIR="$LAMMPS_BUILD/lib/kokkos"
+cmake --build build-plugin-gpu-pjrt -j4
+```
+
+## Run
+
+```bash
+python examples/export_model.py lj examples/lj.lammps-jax.json \
+  --max-atoms 1024 --edges-per-atom 96
+```
+
+The bundle loads as a pair style:
+
+```
+pair_style jax/kk ${pjrt}
+pair_coeff * * examples/lj.lammps-jax.json
+```
+
+```bash
+LAMMPS_PLUGIN_PATH=build-plugin-gpu-pjrt \
+  "$LAMMPS_INSTALL/bin/lmp" -k on g 1 -pk kokkos newton off neigh full -sf kk \
+  -var pjrt /absolute/path/xla_cuda_plugin.so -in examples/in.lj_jax
+```
+
+## CPU
+
+Kokkos OpenMP builds run bundles exported with `platform="cpu"` on the XLA CPU
+PJRT plugin from the `cpu` extra (Linux x86-64):
+
+```bash
+uv pip install -e '.[cpu]'
+python -c 'import lammps_jax_pjrt_cpu as p; print(p.PLUGIN_PATH)'
+```
+
+Build LAMMPS with `-C cmake/presets/kokkos-openmp.cmake -D PKG_PLUGIN=on
+-D BUILD_MPI=on -D Kokkos_ENABLE_IMPL_VIEW_LEGACY=on` and the plugin as above
+against that build, without `CMAKE_CXX_COMPILER`. One MPI rank per core runs
+fastest:
+
+```bash
+python examples/export_model.py eam cuzr.lammps-jax.json \
+  --setfl examples/potentials/CuZr.eam.alloy.gz \
+  --max-atoms 8192 --edges-per-atom 128 --precision float64 --platform cpu
+LAMMPS_PLUGIN_PATH=build-plugin-cpu-pjrt \
+  mpirun -np 4 "$LAMMPS_CPU_INSTALL/bin/lmp" -k on t 1 -sf kk -pk kokkos newton on neigh half \
+  -var pjrt /absolute/path/pjrt_c_api_cpu_plugin.so -in examples/in.eam_cuzr
+```
+
+cuEquivariance and OpenEquivariance kernels are CUDA only.
+
+## Test
+
+```bash
+JAX_PLATFORMS=cpu python -m pytest
+```
+
+Integration tests, which require a GPU and LAMMPS:
+
+```bash
+LAMMPS_BIN=$LAMMPS_INSTALL/bin/lmp \
+PJRT_PLUGIN=/absolute/path/xla_cuda_plugin.so \
+LAMMPS_PLUGIN_PATH=$PWD/build-plugin-gpu-pjrt \
+  python -m pytest tests/test_lammps.py -v
+```
+
+For a host Kokkos LAMMPS and the CPU plugin, add `LAMMPS_JAX_TEST_PLATFORM=cpu`.
+
+## Design notes
+
+Models are exported with `jax.export` into a JSON bundle holding the
+serialized program and the settings the pair style enforces at load: atom and
+edge capacities, cutoff, precision, unit style, and the distribution scheme.
+Exported programs have static shapes, so positions and the edge list are
+padded to capacity and an edge mask marks the live entries. The pair style
+distance-filters pairs to the model cutoff and repacks the edge list every
+step, so edge capacity sizes to the cutoff rather than cutoff plus skin.
+
+The pair style executes the bundle through the PJRT C API, loading the same
+plugin library jax uses on the GPU, so the build never needs XLA. Kokkos kernels
+pack the LAMMPS neighbor list into sender, receiver, and mask arrays on
+device, the program runs on the LAMMPS CUDA stream, and forces come back as
+device buffers, so array data never passes through the host.
+
+Multi-rank runs pick one of two schemes at export. A ghost export widens the
+ghost shell to n_hops cutoffs and stays communication-free inside the
+program; a comm export keeps the one-cutoff shell and exchanges per-layer
+features through LAMMPS forward and reverse communication, called from inside
+the program as an FFI callback. Exchanges are linear JAX primitives whose
+transposes swap forward and reverse, so models may call either direction and
+every differentiation mode traces through them. Communicating half-edge
+exports pack each pair on one rank, matching the native half neighbor list,
+and reverse-communicate ghost density partials into their owners. Exchanged
+rows move in place on device under the Kokkos brick and tiled comm styles and
+stage through pinned host memory otherwise.
+
+cuEquivariance and OpenEquivariance kernels stay in the exported program as
+custom call targets, resolved at run time from LAMMPS_JAX_FFI_HANDLERS;
+contrib/ffi-replay handles libraries whose compiled kernels only exist inside
+the exporting process. NequIP checkpoints export through the uv-runnable
+scripts in examples/torch, which trace the torch blocks via torchax into
+the serialized program; the bundles load and run like any native export.
+
+## Citation
+
+If you use `lammps-jax` or any algorithms implemented here, please cite
+the archived Zenodo release: https://doi.org/10.5281/zenodo.20838694
+
+## Acknowledgements
+
+Thanks to @mitkotak for API discussions and @wcwitt for guidance on
+distributed inference. See
+[lammps/lammps#4691](https://github.com/lammps/lammps/pull/4691) for the
+original discussion on running JAX models in LAMMPS, and
+[openmm-jax](https://github.com/atomicarchitects/openmm-jax) for the OpenMM
+counterpart.
+
+## License
+
+MIT. See `LICENSE`.

lammps_jax-0.2.0/pyproject.toml

@@ -0,0 +1,65 @@
+[build-system]
+requires = ["setuptools>=77", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "lammps-jax"
+version = "0.2.0"
+description = "CUDA-first LAMMPS/Kokkos plugin support for exported JAX models"
+readme = "README.md"
+license = "MIT"
+requires-python = ">=3.10"
+# Comm bundles need the plugin's FFI registration from 0.10.0; the 0.10.2
+# plugin slows GPU scatter-add about 4x, so run 0.10.1.
+dependencies = [
+  "jax>=0.8",
+  "jaxlib>=0.8,!=0.10.2",
+]
+
+[project.urls]
+Repository = "https://github.com/abhijeetgangan/lammps-jax"
+
+[project.optional-dependencies]
+# The XLA CPU PJRT plugin for host Kokkos builds (pjrt-cpu/); jax ships none.
+cpu = [
+  "lammps-jax-pjrt-cpu>=0.11.2; sys_platform == 'linux' and platform_machine == 'x86_64'",
+]
+dist = [
+  "ase>=3.22",
+  "e3nn-jax>=0.20",
+  # jax 0.11.2 segfaults compiling ghost_exchange.redistribute on CPU; 0.11.1 passes.
+  "jax!=0.11.2",
+  "jax-md>=0.2.27",
+  "nequix>=0.4",
+]
+test = [
+  "pytest>=7",
+  "numpy>=1.24",
+]
+
+[tool.setuptools]
+package-dir = {"" = "python"}
+
+[tool.setuptools.packages.find]
+where = ["python"]
+
+[tool.pytest.ini_options]
+pythonpath = ["python"]
+testpaths = ["tests"]
+
+[tool.basedpyright]
+venvPath = "."
+venv = ".venv"
+extraPaths = ["python", "tests"]
+
+[tool.ty.environment]
+python = "./.venv"
+root = ["python", "tests"]
+
+[[tool.ty.overrides]]
+include = ["tests/**"]
+
+[tool.ty.overrides.rules]
+invalid-argument-type = "ignore"
+invalid-assignment = "ignore"
+unsupported-operator = "ignore"

lammps_jax-0.2.0/python/lammps_jax/__init__.py

@@ -0,0 +1 @@
+"""LAMMPS-JAX package."""