chex>=0.0.6
numpy>=1.18.0
dataclasses==0.8; python_version < "3.7"
