Specify that mask must be static in masked wrapper docstring.
Specify that mask must be static in masked wrapper docstring.

The mask in the masked wrapper must be static to be able to use jax transforms
such as jax.jit on the init and update functions of the optax gradient
transformation. This is now stated in the docstring of the wrapper.

PiperOrigin-RevId: 430161629

deepmind/optax,deepmind/optax