dm-haiku
frozendict
jax
jaxlib
optax
tensorflow-federated
