absl-py
flax
google-benchmark
jax
ml_collections
numpy
optax