Method step

python/tests/test_compile.py:666–668

(w)

            return w @ w1 + w2 @ w2

        def step(w):
            out, grad = mx.vjp(fun, (w,), (mx.array([[1.0, 1.0], [1.0, 1.0]]),))
            return out[0], grad[0]

        w = mx.zeros((2, 2))
        mx.eval(w)