Fix grad in place updates (#2899)

2025-12-16 01:49:05 +08:00 · 2025-12-11 14:44:58 -08:00
parent 012fb220a1
commit 81dfe5f137
3 changed files with 43 additions and 15 deletions
--- a/python/tests/test_vmap.py
+++ b/python/tests/test_vmap.py
@@ -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)