CommonLoopUtils
CommonLoopUtils copied to clipboard
google
CLU lets you write beautiful training loops in JAX.
![copybara-service[bot] avatar](https://avatars.githubusercontent.com/in/44061?v=4 "copybara-service[bot] avatar"){kind=avatar}
Create separate module for TF-compatible `parameter_overview`.
metrics: Use Array and ArrayLike types thoughout Currently the inputs to `from_model_output` are not typed. However, these functions cannot accept arbitrary inputs, they need to be a value convertable to...
Make Metrics and Collection class inherit from struct.PyTreeNode, so that we don't have to decorate them with struct.dataclass anymore.
