jax
numpy
numba
h5py
scipy
sympy
flax[all]
tomlkit
optax
ase
pytest

[cuda]
jax[cuda12]

[deephp]
cffi
pydlpack
pycuda

[e3nn]
e3nn-jax
