CLU lets you write beautiful training loops in JAX.
Changelog:
Adds metrics module and some minor changes.
Current list of modules:
clu.checkpoint
clu.deterministic_training
clu.metric_writers
clu.periodic_actions
clu.platform
clu.profiler