jax>=0.9
numpy
orbax-checkpoint
termcolor

[cuda12]
jax[cuda12]>=0.9

[cuda12-local]
jax[cuda12-local]>=0.9

[cuda13]
jax[cuda13]>=0.9

[cuda13-local]
jax[cuda13-local]>=0.9

[dev]
build
twine
ruff
pytest

[docs]
mkdocs
mkdocs-material
mkdocstrings[python]
mkdocs-autolinks-plugin
pymdown-extensions

[io]
pyvista
matplotlib
Rtree
trimesh
h5py

[rocm]
jax[rocm7-local]>=0.9

[tpu]
jax[tpu]>=0.9
