v0.2.11

• Use jax.tree_util.* functions to avoid warnings.
• Full Changelog: v0.2.10...v0.2.11

v0.2.10

• ScaleBySchedule: correct the learning rate at initialization.

Full Changelog: v0.2.8...v0.2.10

v0.2.8

Introduce strict and non-strict modes:

• Strict mode requires pytrees to be matched.
• Non-strict mode only requires tree leaves to be matched.

v0.2.7

• Make testcases compatible with pax3>=0.5.0.

v0.2.6

• Simplify the implementation by considering gradient transformations as StateModule.

v0.2.5

Introduce the class method init. It makes things easier to understand because we explicitly say "initialize the optimizer".

Implicit:

optimizer = opax.adam(1e-3)(net.parameters())

Explicit:

optimizer = opax.adam(1e-3).init(net.parameters())

v0.2.4

• Import add_decayed_weights.
• Add __all__ to __init__.py file.

v0.2.3

• Use the newest PAX v0.4.0.
• Remove step method.
• Add transform_gradients, apply_updates and apply_gradients utility (pure) functions.

v0.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.

v0.1.0

First official relase:

• support adam, adamw, rmsprop and sgd optimizers.
• support lr schedule.