kd.optim.partial_updates

kauldron.optim.partial_updates(

optimizer: optax._src.base.GradientTransformation,

mask: Any | collections.abc.Callable[[Any], Any],

) → optax._src.base.GradientTransformation

Applies the optimizer to a subset of the parameters.

Parameters:

• optimizer – The optimizer to use.

• mask – A tree or callable returning a tree of bools to apply the optimizer to.

Returns:

The wrapped optimizer.