jax
jaxtyping>=0.3.3

[tests]
pytest
