absl-py>=2.3.0
einshape
jax>=0.11.0
jaxlib>=0.11.0
jaxtyping>=0.3
pydantic>=2.11.0
qwix>=0.1.2
tqdm
typeguard==2.13.3
immutabledict
tensorboardx
xprof>=2.23.1

[bench]
google-benchmark>=1.9.0

[cuda]
jax[cuda13]>=0.11.0
nvidia-cudnn-cu13>=9.23.0.39
cuequivariance-jax>=0.10.0
cuequivariance-ops-jax-cu13>=0.10.0
triton

[test]
chex>=0.1.91
flatbuffers
flax
pytest-xdist
pytest
hypothesis

[tpu]
jax[tpu]>=0.11.0
