jax
jaxlib

[losses_utils]
dm-haiku
