Code example for stop_gradient()

[RahulBhalley]

Could someone please help me with a code example how to use stop_gradient? In MLX docs it's said we don't need detach as in PyTorch but I wonder if stop_gradient fulfils the same purpose.

[awni]

def fun(x):
  return mx.exp(x) + mx.stop_gradient(mx.exp(x))

print(mx.grad(fun)(mx.array(1.0)))

Gives array(2.71828, dtype=float32).

So there you would only get the gradient through the first mx.exp.

Compare to:

def fun(x):
  return mx.exp(x) + mx.exp(x)

print(mx.grad(fun)(mx.array(1.0)))

Which gives you twice the result of the first (grad through both paths). Gives array(5.43656, dtype=float32).