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?