pip install --no-deps holisticai jax jaxlib ml-dtypes optax chex toolz etils flax msgpack rich
# opt-einsum absl-py importlib-metadata