jax
jaxlib
numpy

[dev]
pytest
equinox
optax
tqdm
stable-gymnax
diffrax
