# JAX MD for differentiable molecular dynamics

> JAX MD implements molecular dynamics as functional Python code over JAX arrays, targeting research that needs accelerated execution or derivatives through simulation code. It offers a flexible modeling environment, but its research status and documented sharp edges make it unsuitable as a drop-in replacement for specialized production packages.

**jax-md/jax-md** — Differentiable, Hardware Accelerated, Molecular Dynamics

- Repository: https://github.com/jax-md/jax-md
- Stars: 1,473 · Forks: 253
- Language: Jupyter Notebook
- License: Apache-2.0
- Published: 2026-10-16 · Updated: 2026-10-16 · Language: en
- Canonical page: https://hysenlabs.com/projects/jax-md-jax-md

## JAX MD's functional approach to particle simulation

JAX MD targets condensed matter researchers, machine learning researchers, and developers who want molecular dynamics expressed as differentiable Python. The README contrasts it with specialized systems such as HOOMD Blue and LAMMPS, which it describes as large, duplicated codebases that maintain separate CPU and GPU paths. JAX MD instead places simulation logic on top of JAX, allowing Python functions to be compiled through XLA and executed on CPU, GPU, or TPU. This is useful when one experiment needs both simulation and machine learning components built around the same array ecosystem.

The library also aims to be end-to-end differentiable. Forces can require derivatives of energy, while JAX can derive Python functions. That combination makes JAX MD relevant to experiments involving optimization, learned potentials, implicit differentiation, and other tasks where gradients through a simulation may matter. The repository includes notebooks for custom potentials, meta optimization, implicit differentiation, neural network potentials, and athermal elasticity, which gives the feature set more context than a general-purpose simulator description alone.

Its intended audience is therefore narrower than the label “molecular dynamics” might suggest. JAX MD is designed for research software that benefits from composable transformations and automatic differentiation. It is not presented as a complete replacement for the long-established feature set and operational conventions of dedicated simulation packages.

## Array states and transformed functions in JAX MD

Based on the README and repository layout, JAX MD is organized as a functional, data-driven library. State is held in arrays or tuples of arrays, and library functions transform one state into another. Particle positions, velocities, forces, and related quantities can therefore flow through ordinary function composition rather than being managed through a large mutable simulation object. The `jax_md/` package contains the library, while `examples/` and `notebooks/` show intended simulations and research workflows.

Execution delegates much of the hardware behavior to JAX and XLA. According to the README, code written for JAX can be compiled and run across CPU, GPU, and TPU. JAX's differentiation machinery supplies derivatives without requiring the library to maintain separate derivative implementations for every simulation function. This is the central architectural trade-off: concise functional code can be transformed and differentiated, but users must work within JAX's execution and numerical model.

The repository supports several simulation-oriented examples, including `nve_simulation.py`, `nvt_simulation.py`, `npt_simulation.py`, neighbor-list examples, minimization, neural network potentials, symmetric dynamics, and external examples named for `reaxff` and `uma`. The examples directory also contains data, models, units, and custom neighbor-list code. That spread suggests JAX MD is intended to cover multiple state ensembles and potential implementations, although the README does not claim parity with every feature found in specialized simulation software.

## Installing JAX MD with uv or pip

For an installation-free first look, the README recommends the JAX MD Cookbook and Minimization material hosted in Google Colab. The cookbook is the better starting point when the goal is to inspect several features, while the README specifically recommends the Minimization example as a simple introduction. These notebooks use the repository's `google/jax-md` notebook links and do not require a local package installation.

For local development with uv, clone the repository, enter its directory, and synchronize the default project environment. The documented sequence is:

```bash
git clone https://github.com/jax-md/jax-md
cd jax-md
uv sync --no-default-groups
```

After this sequence completes, the repository says self-contained Python examples can be run locally with JAX MD installed. It lists `examples/fire_minimization.py` as a minimization example, along with NVE, NVT, NPT, neighbor-list, and neural network potential scripts. The README does not give a dedicated invocation command for those scripts, so it does not establish a required launcher or extra flag.

The project metadata sets `requires-python = ">=3.10"`, and the current package version in `pyproject.toml` is `0.2.29`. A pip installation uses the published package:

```bash
python -m pip install jax-md --upgrade
```

For a local source checkout, the documented editable installation is:

```bash
python -m pip install -e .
```

## First simulation setup and the x64 requirement

JAX MD's test suite gives a concrete indication that 64-bit precision matters for its simulations. The README says tests run in double precision and requires `JAX_ENABLE_X64=1` before pytest. It also says the corresponding JAX flag must be set before importing `jax_md`. The documented Python setup is:

```python
import jax
jax.config.update("jax_enable_x64", True)
```

This code enables JAX's 64-bit mode. The README does not show a complete import and simulation program, so it would be misleading to invent initialization parameters, particle arrays, integration steps, or output values. The concrete local starting point is `examples/fire_minimization.py`, which the repository identifies as a self-contained Fire minimization script.

GPU support has a separate boundary. JAX MD does not bundle its own accelerator setup procedure. The README directs users to JAX's GPU installation instructions, meaning JAX and accelerator configuration form part of the deployment decision. A working CPU installation therefore does not establish that GPU or TPU execution is configured.

Dependency compatibility also requires attention. `pyproject.toml` accepts `jax>=0.5.0` and `jaxlib>=0.5.0`, but excludes `0.10.2` and `0.11.1`. The metadata says `0.11.1` drops `jax.experimental.hijax.MutableHiType`, which `flax 0.12.8` imports at module scope, while `0.10.2` is excluded in response to a JAX issue. Those constraints are concrete signs that JAX MD tracks a changing JAX ecosystem.

## Research flexibility against production specialization

JAX MD's main limitation is stated plainly by the README: it is a research project under development, with expected “sharp edges” and possible API-breaking changes. This warning matters because differentiating or JIT-compiling simulation code introduces constraints that cannot be judged from installation alone. A functional model may be concise, but the complete research still depends on valid JAX arrays, transformation behavior, numerical precision, potential definitions, and the stability of each integration scheme.

The alternative approach is to use a specialized package such as HOOMD Blue or LAMMPS. The README identifies these as alternatives that organize molecular dynamics as substantial, specialized codebases. That design can suit established simulation workflows, but the README does not provide a feature-by-feature comparison or claim that either package supports end-to-end differentiation through the same Python functional model. The defensible distinction is architectural: JAX MD treats simulation as differentiable JAX transformations, while specialized packages center their design on conventional molecular dynamics execution.

JAX MD can also be the wrong tool when the experiment does not need derivatives or a JAX-centered workflow. The added dependency on JAX's execution model, x64 configuration, and version constraints would not be justified merely because another package is written in Python. Conversely, it is a credible candidate when a neural network potential, parameter optimization, or custom simulation function must participate in a differentiable or compiled computation.

## JAX MD development, tests, and upgrade exposure

The repository is not archived, and its last push was on 2026-08-18. Recent releases include `jax-md-v0.2.29` on 2026-06-10, `jax-md-v0.2.28` on 2026-03-22, and `jax-md-v0.2.27` on 2025-11-07. These facts show recent project activity, but the README still characterizes JAX MD as under active development and warns about sharp edges and API changes. Current repository activity does not remove that compatibility risk.

The test workflow is documented in enough detail to form an upgrade checkpoint. After installing the testing group, the full suite can run directly through pytest:

```bash
python -m pip install -e . --group testing
JAX_ENABLE_X64=1 python -m pytest
```

A contributor can instead target one suite using the documented filename pattern:

```bash
JAX_ENABLE_X64=1 python -m pytest tests/<suite>_test.py
```

The literal `<suite>` is a placeholder in the repository documentation, not a test filename that can be copied literally without substitution. The tests live in `tests/`, and the README says to run suites affected by a change. For project users, this matters most when upgrading JAX, `jaxlib`, Flax, or a potential model. Double-precision tests provide one documented compatibility check, but they are not a substitute for validating the user's own simulation and accelerator configuration.

JAX MD uses the `Apache-2.0` license, with `LICENSE` and `LICENSE_SHORT` declared in package metadata. That identifies the project's stated license, but it does not by itself settle questions about dependencies, models, data, or a particular deployment.

## Conclusion

JAX MD fits researchers who need to differentiate through molecular dynamics, prototype custom potentials, or compose simulation code with JAX-based machine learning. It is a poorer choice for teams that require the specialized feature set and predictable production workflow associated with packages such as HOOMD Blue or LAMMPS. Before adoption, verify the required example in `examples/`, confirm accelerator setup through JAX's own instructions, check the 64-bit precision path, and review compatibility with the intended JAX version because `jax` and `jaxlib` have explicit exclusions and API changes remain possible.

## FAQ

### What is JAX MD used for?

JAX MD is a Python molecular dynamics library for research that benefits from JAX execution across CPU, GPU, or TPU and from automatic differentiation through simulation code. Its examples cover minimization, several thermodynamic ensembles, neighbor lists, custom potentials, and neural network potentials.

### How do you install JAX MD?

The README recommends `uv sync --no-default-groups` for a cloned local repository. For the published package, it documents `python -m pip install jax-md --upgrade`.

### Does JAX MD require 64-bit precision and GPU setup?

The README says 64-bit mode must be enabled before importing `jax_md` and shows `jax.config.update("jax_enable_x64", True)`. GPU support is separate, and users must follow JAX's GPU installation instructions.

## Sources

- [Issues](https://github.com/jax-md/jax-md/issues)
- [jax-md/jax-md on GitHub](https://github.com/jax-md/jax-md)
- [License: Apache-2.0](https://github.com/jax-md/jax-md/blob/main/LICENSE)
- [README](https://github.com/jax-md/jax-md/blob/main/README.md)
- [Releases](https://github.com/jax-md/jax-md/releases)

---

Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/jax-md-jax-md