absl-py
cached_property
dataclasses
flax
jax
jaxlib
numpy
pyyaml
tensorflow
tensorstore>=0.1.20

[test]
pytest
