einops>=0.5.0
jax>=0.4.14

[dev]
pre-commit
pytest
ruff
