jax
jaxlib
flax
einops
tensorflow
numpy
optax
tqdm
