numpy
scipy

[jax]
jax
optax
