jax
dm-haiku
optax
sympy
numpy
