jax
dm-haiku
flax
optax
sympy
numpy
tqdm
