jax_tools Repo that contains personal scripts that I use to build Machine Learning algorithms on top of JAX, HAIKU, OPTAX...