numpy
scipy
array_api_compat

[jax]
jax
optax
