chex
dm-haiku
flax
# jax
jaxopt
matplotlib
ml-collections
optax
plotly
wandb