kd.optim#
Optimizers etc.
Symbols#
Function#
Add (params - init_params) scaled by |
|
Create a mask which selects all nodes except the ones matching the pattern. |
|
Wraps optax.named_chain and allows passing transformations as kwargs. |
|
Applies the optimizer to a subset of the parameters. |
|
Create a mask which selects only the sub-pytree matching the pattern. |