jax
jmp
opax

[test]
chex
dm-haiku
fire
mypy
opax
pytest
tensorflow_datasets
tensorflow
tqdm
