jax
dm-haiku
optax
sympy
