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

[stockeh]

@awni It may be helpful to have a no_grad decorator to simplify the following:

import functools

import mlx.core as mx


def no_grad(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        out = fn(*args, **kwargs)
        if isinstance(out, mx.array):
            return mx.stop_gradient(out)
        elif isinstance(out, (tuple, list)):
            return type(out)(
                mx.stop_gradient(x) if isinstance(x, mx.array) else x for x in out
            )
        else:
            return out

    return wrapper


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

print(mx.grad(fun)(mx.array(1.0)))
# array(0, dtype=float32)

[awni]

Just curious what did you have in mind to use that for? The no_grad idiom is not very common in MLX because it doesn't by default keep the tape around for comptuing gradients.

Like mx.eval(fun(x)) is just as memory efficient as mx.eval(mx.stop_gradient(fun(x)))

[stockeh]

I'm reimplementing TRM (see Fig 3, https://arxiv.org/abs/2510.04871) in MLX and looking to have the no_grad components within the model's call, similar to the original implementation here. I was wanting to have a context or decorator to achieve this. Am I thinking of things backwards with best practices or do you have any recommendations?

[awni]

I'm reimplementing TRM

Cool!!

[awni]

I see. I think your decorator is fine for that use case if you split the no grad part out into it's own function. But it's also quite simple to do something like the following (which is equivalent):

        # H_cycles-1 without grad
        for _H_step in range(self.config.H_cycles-1):
            for _L_step in range(self.config.L_cycles):
                z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info)
            z_H = self.L_level(z_H, z_L, **seq_info)

        # Add that
        z_H, z_L = mx.stop_gradient(z_H), mx.stop_gradient(z_L)

        # 1 with grad
        for _L_step in range(self.config.L_cycles):
            z_L = self.L_level(z_L, z_H + input_embeddings, **seq_info)
        z_H = self.L_level(z_H, z_L, **seq_info)

[stockeh]

Simple indeed, that totally makes sense! Thank you.