jax
jaxlib
numpy
optax

[dev]
pytest
ruff
mypy
pre-commit

[torch]
torch
