torch==1.9.0
einops==0.3.0
dm-haiku
"jax[cpu]"
