Improve the documentation and discoverability of `set_to_zero`.
Improve the documentation and discoverability of `set_to_zero`.

This commit improves the docstring of `set_to_zero` to

* Increase discoverability by including more expressions users might search for
  ("freeze" and "keep fixed").
* Explain the performance consideration when using `set_to_zero`.

This addresses issue [#290](https://github.com/deepmind/optax/issues/290).

PiperOrigin-RevId: 426350880

deepmind/optax,deepmind/optax