- Import `add_decayed_weights`. - Add `__all__` to `__init__.py` file.
0.2.3
- Use the newest PAX v0.4.0. - Remove `step` method. - Add `transform_gradients`, `apply_updates` and `apply_gradients` utility (pure) functions.
0.1.7
* Check if param's tree structure matchs gradient's tree structure for better error messages. * Support `all_finite` predicate when updating parameters. This is useful when models are trained with mixed precision. * Support `flatten` mode to speed up `tree_flatten` and `tree_unflatten`.
0.1.0
First official relase: - support `adam`, `adamw`, `rmsprop` and `sgd` optimizers. - support lr schedule.