jax
optimistix
optax
