jax-md has 520K monthly downloads and 1K GitHub stars. Version 0.2.29 supports Python 3.10+, first published in 2019.
It is actively maintained under a permissive license, last released in Jun 2026.
It is built on top of absl-py, e3nn-jax and einops. It is used by 8 tracked packages, including proxide, mlip and mlipaudit.
Package Insights
((week_daily_avg - month_daily_avg) / month_daily_avg) * 100Weekly Downloads
GitHub Stars
Downloads by OS
Python Versions
Top Countries
Dependencies
- absl-py
- e3nn-jax
- einops
- flax
- frozendict
- jax >=0.5.0
- jaxlib >=0.5.0
- jraph
- ml-collections
- numpy
- optax
- pymatgen >=2021.3.4
3 optional dependencies
- openmm[amber]
- openmmforcefields[amber]
- parmed[amber]