add freezing parameters examples [Addresses Issue #296]

[pranavagrawaI]

What’s Changed

• New example notebook: examples/freezing_parameters.ipynb
  • 3 concise recipes showing how to freeze parameters in Optax:
    1. Simple masking with optax.masked + set_to_zero
    2. Group freezing via optax.multi_transform + set_to_zero
    3. Dynamic scheduling (swap freeze/train mid‑run)
  • Self‑contained cells, minimal JAX arrays, and inline sanity checks (assert/print) with graphical visualizations.

[vroulet]

It looks great.
Fixes #296
Use optax.partition instead of optax.multi_transform.

Ideally, we could add the following shortcuts in the utils file (with docstring and tests associated). Let me know if you are willing to do this or not.

def freeze(mask):
  return optax.masked(optax.set_to_zero(), mask)

def partial_updates(optimizer, updatable_params):
  # say updatable_params is a boolean pytree
  # we'll need to convert it in a label_params for optax.partition
  label_params = ... 
  return optax.partition(
    {
        'train': optimizer,
        'freeze': optax.set_to_zero(),
    },
    label_params,
)

[vroulet]

optax.partition (not optix.multi_transform)

[pranavagrawaI]

Thanks for the feedback, I am glad it looks good.

1. Switched to optax.partition.
All examples now use partition instead of multi_transform.

2. I would be glad to implement the shortcuts and along with tests and doc strings, here are a few helper function ideas (guidance appreciated).

• freeze(mask)
  A simple wrapper around optax.masked(optax.set_to_zero(), mask) for convenience.
• partial_updates(optimizer, mask)
  Builds on partition to apply your chosen optimizer to True entries in mask and freeze the rest.
• freeze_layers(layer_names, params)
  Generates a boolean mask that freezes all parameters belonging to specified layers or collections.
• mask_by_pattern(regex, params)
  Creates a mask by matching parameter paths against a regular expression (e.g. freeze all “bias” or “batch_norm” terms).
• scheduled_freeze(mask_fn, schedule)
  A higher‑order transform that applies different masks over training steps (e.g. unfreeze after N steps).

Any thoughts on which of these would be most valuable, or would the first 2 you mentioned suffice?

3. Doc update for set_to_zero.
At the moment, the documentation for set_to_zero says that multi_transform or masked can be combined with set_to_zero to freeze parameters. I can update the documentation to point users toward partition as the preferred freeze pattern (In a separate pull request or here), replacing the current multi_transform recommendation.

Let me know what helper functions you’d like prioritized or if you have other ideas!

[vroulet]

The example looks great. If you could indeed also update set_to_zero doc that would be perfect. Then we'll merge this one.

You may then do another PR for the shortcuts. They can be put in the utils file of the _src folder.
Freeze and partial_updates would be great. Happy to look at the PR once it's ready. Let me know if the overall signatures I sent you were not clear.

For freeze_layers and mask_by_pattern,

• I think these could be achieved by the tree_set functionality in tree_utils.
• masked accepts callables so they can be achieved easily by a callable.
  I think here having a longer docstring for optax.masked would be the best but this can be for another PR and is not urgent.

I don't think we need a scheduled_freeze. Your example is enough for people to catch how to do it.

Thank you very much @pranavagrawaI

[pranavagrawaI]

I have changed multi_transform to partition in the set to zero doc string. I think this should be ready to merge.

[vroulet]

Thanks! Let's merge this