What does optax.mask do?

I don’t understand what optax.mask does.

I would expect that the masked optimizer

optax.mask(optax.sgd(0.1), {"Dense": True, "bias": False})

would only apply the optimisation to sub-leafs of Dense and not optimise sub-leaves of bias. Which means that the masked gradient should match the sgd one for Dense and be zero for bias.

However it seems to me that the masked updates are correct for sub-leafs of Dense (so where the mask is True, but they are the identity where the mask is False.

Is this intended behaviour? it seems rather strange to me. I was trying to update only a subsets of the weights of my model but this was not working

MWE:

import jax.numpy as jnp
import jax
import optax

pars = {"Dense": {"kernel": jnp.zeros((2,3)), "bias": jnp.zeros((3))}, "bias":jnp.zeros(2)}
grad = jax.tree_map(jnp.ones_like, pars)

op = optax.masked(optax.sgd(0.1), {"Dense": True, "bias": False})

op_state = op.init(pars)

masked_updates, new_op_state = op.update(grad, op_state,  pars)
>>> masked_updates
{'Dense': {'bias': DeviceArray([-0.1, -0.1, -0.1], dtype=float32), 'kernel': DeviceArray([[-0.1, -0.1, -0.1],
             [-0.1, -0.1, -0.1]], dtype=float32)}, 'bias': DeviceArray([1., 1.], dtype=float32)}

[PhilipVinc]

I think it does, because Dead Code Elimination will see that they are not used anywhere. (zero_grad calls zeros_like)

[PhilipVinc]

@pharringtonp19 sorry i was on holiday. @n2cholas I eventually had figured this out. I admit, this behaviour makes complete sense within the framework of gradient transformation of optax, however it took me a while to figure it out.

the current docstrings states Mask updates so only a subset of them are computed. . In optax view, that means that only a subset of the gradient are transformed, and the rest arent. But if one isn’t familiar with optax, i would immediately think that only some gradients are computed, and the rest are zeroed out or whatever.

I think the docstring could be modified to state something like:

  """Mask updates so only a subset of them are transformed, while the rest are 
passed-through unchanged.
"""