einops>=0.3
jax>=0.2.10
