Method f

python/tests/test_array.py:493–495

(x)

        # check that `[mx.array, ...]` retains the `mx.array` in the graph
        def f(x):
            y = mx.array([x, mx.array([2.0])])
            return (2 * y).sum()

        x = mx.array([1.0])
        dfdx = mx.grad(f)