Fix grad in place updates (#2899)

2025-12-11 14:44:58 -08:00
parent 012fb220a1
commit 81dfe5f137
3 changed files with 43 additions and 15 deletions
@@ -744,7 +744,6 @@ class TestVmap(mlx_tests.MLXTestCase):
            return Vector([t[0] + 10, t[1] * 10])

        x = State(mx.array(1), mx.array(2))
-        print(f"{transform(x)=}")

        vmap_transform = mx.vmap(transform)
        vmap_transform_tuple = mx.vmap(transform_tuple)