jax>=0.9
numpy
orbax-checkpoint

[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]
zensical
mkdocstrings-python

[io]
pyvista
matplotlib
Rtree
trimesh
h5py

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

[tpu]
jax[tpu]>=0.9
