jax>=0.4
equinox>=0.11
einops
