flax
jax
numpy
