
[all]
torch>=2.2.0
jax>=0.4.34

[dev]
pytest>=7.0

[jax-cpu]
jax[cpu]>=0.4.34

[jax-gpu]
jax[cuda12]>=0.4.34

[torch]
torch>=2.2.0
