Can you stop gradients for part of a tensor?

[cianciosa]

I'm trying to gradients off for a subset of a tensor. Take the example

import mlx.core as mx

a = mx.array([1,2,3])
b = mx.where(a == 3, mx.stop_gradient(a), a)

def f(x):
    return mx.sum(x)

grad_f = mx.grad(f)

print(grad_f(b))
print(grad_f(a))

In this, I'm trying to turn off gradients for the 3rd element of a. When I evaluate this, I get the result

array([1, 1, 1], dtype=int32)
array([1, 1, 1], dtype=int32)

I would have expected this to produce

array([1, 1, 0], dtype=int32)
array([1, 1, 1], dtype=int32)

Can you turn off gradients for part of a tensor?

[aicayzer]

fwiw mx.grad(f) just gives you df/dx evaluated at whatever you pass in — it doesn't trace back through how that value got constructed. So grad_f(b) returns the gradient of sum at b, which is [1, 1, 1] regardless of how you built b.

To actually test stop_gradient through where(), the where() needs to be inside the function you're differentiating, and you take the gradient with respect to the original a:

import mlx.core as mx

a = mx.array([1.0, 2.0, 3.0])

def f(x):
    b = mx.where(x == 3, mx.stop_gradient(x), x)
    return mx.sum(b)

grad_f = mx.grad(f)
print(grad_f(a))   # array([1, 1, 0], dtype=float32)

That gives the [1, 1, 0] you were after.

(Side note: mx.grad needs float inputs — your mx.array([1,2,3]) is int32. That's also why your printout had int gradients. Quietly works in your example because sum is linear, but it'll bite you if f has any actual computation.)

[yentur]

The fix in the comment above is right, and it holds on mlx 0.32.0: the where has to be inside the function you differentiate. mx.grad traces the function it is handed and treats its arguments as leaves, so the graph you built into b beforehand is not part of that trace.

Using a cube instead of a plain sum so the live entries are not all 1 and the masking is actually visible:

import mlx.core as mx

a = mx.array([1.0, 2.0, 3.0])
mask = mx.array([False, False, True])

def g(x):
    y = mx.where(mask, mx.stop_gradient(x), x)
    return mx.sum(y ** 3)

print(mx.grad(g)(a))   # array([3, 12, 0], dtype=float32)
print(3 * a ** 2)      # array([3, 12, 27], dtype=float32)

The forward pass is untouched (mx.all(mx.where(mask, mx.stop_gradient(a), a) == a) is True), and I get the same [3, 12, 0] out of m * mx.stop_gradient(x) + (1 - m) * x, out of mx.concatenate([x[:2], mx.stop_gradient(x[2:])]) when the frozen part is contiguous, and out of mx.compile(mx.grad(g)).

On the int32 arrays in your snippet: the gradient is cast back to the input dtype, so it does not merely look odd, it can come out zero.

print(mx.grad(lambda x: mx.sum(mx.sqrt(x)))(mx.array([1, 2, 3])))
# array([0, 0, 0], dtype=int32)
print(mx.grad(lambda x: mx.sum(mx.sqrt(x)))(mx.array([1.0, 2.0, 3.0])))
# array([0.5, 0.353553, 0.288675], dtype=float32)

One thing that bit me putting this in an actual training loop: a zero gradient is not the same as a frozen parameter for every optimizer. AdamW scales the parameter before handing it to Adam, whatever the gradient is (python/mlx/optimizers/optimizers.py:586-588):

return super().apply_single(
    gradient, parameter * (1 - lr * self.weight_decay), state
)

20 steps on nn.Linear(3, 4, bias=False) with rows 2 and 3 masked, printing max |grad| and max |w - w0| per row:

SGD(momentum=0.9)
  grad per row : array([0.125087, 0.0945172, 0, 0], dtype=float32)
  |dW| per row : array([0.648379, 0.841788, 0, 0], dtype=float32)

Adam
  grad per row : array([0.100603, 0.0849389, 0, 0], dtype=float32)
  |dW| per row : array([0.6928, 0.610856, 0, 0], dtype=float32)

AdamW(weight_decay=0.1)
  grad per row : array([0.0996453, 0.0787354, 0, 0], dtype=float32)
  |dW| per row : array([0.683424, 0.626188, 0.0763116, 0.104929], dtype=float32)

The masked rows have exactly zero gradient in all three, but under AdamW they still drift. So if what you are freezing is part of a weight rather than part of an intermediate, use weight_decay=0 or write the saved values back after opt.update.