flax
jax
jaxlib

[:python_version < "3.7.0"]
dataclasses

[testing]
pytest
pytest-cov
