Stopping gradient through part of an array

[SNMS95]

If I have a function whose input is jax array, can I stop the gradient's w.r.t a few selected elements?
Ideally, it would be something like this.

import jax.numpy as np
import jax

def f(x):
  y = x[0]
  x = jax.lax.stop_gradient(y)
  return np.mean(np.sin(x))

jax.grad(f)(np.array([3., 1.]))

This should return [0.0, 0.27] but instead it gives [0.0, 0.0].

P.s. I am trying to do an operation like convolution on a 3d array. However, a few pixels in the middle are masked and have to be kept to a constant value. Ideally, the 3d colvolution should only affect the non-masked pixels, but I couldn't do that, so I am trying this approach where I simply cut off the gradients [so that they can never change]

[ASEM000]

I use the following pattern - not sure if it is the best approach though -

import jax.numpy as np
import jax


def f(x):
    return np.mean(np.sin(x))


def grad_with_mask(func, mask=None):
    def grad(*args, **kwargs):
        out = jax.grad(func)(*args, **kwargs)
        if mask is None:
            return out
        return jax.tree_map(lambda node, mask: np.where(mask, 0, node), out, mask)

    return grad


print(jax.grad(f)(np.array([3.0, 1.0])))
# [-0.49499625  0.27015114]
print(grad_with_mask(f, mask=np.array([True, False]))(np.array([3.0, 1.0])))
# [0.         0.27015114]

It can be used with pytrees too

[SNMS95]

Thanks @ASEM000 for the reply. This will be useful when you want to overwrite gradients!

However, this does not prevent higher-order gradients right nor does it prevent the gradient calculation right?
All we are doing is calculating the gradient and then setting some to zero.

[ASEM000]

[so that they can never change]

You are right. preventing gradient calculation in parts of an array is different from my proposal here that masks their gradients with zeros, even if they achieve the same end goal.

[anh-tong]

Does something like this solve your problem?

def f(x):
    x_grad = x.at[0].set(0.)
    x_no_grad = jnp.zeros(x.shape).at[0].set(x[0])
    x = x_grad + jax.lax.stop_gradient(x_no_grad)
    return jnp.mean(jnp.sin(x))