jax>=0.4
jaxlib>=0.4
flax>=0.8

[test]
pytest
