jax
optimistix
optax

[dev]
build
pytest
