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.)