From ef0ec4663aab2ac9e448be3d40e52f1d55ae8116 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fukan=20Veziro=C4=9Flu?= Date: Wed, 8 Jul 2026 16:00:46 +0300 Subject: [PATCH 1/5] Make index autodiff errors explicitly recommend stop_gradient --- mlx/primitives.cpp | 30 ++++++++------ python/tests/test_autograd.py | 77 +++++++++++++++++++++++++++++++++++ python/tests/test_blas.py | 18 ++++++++ 3 files changed, 113 insertions(+), 12 deletions(-) diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index 6ac88eef6c..9c96cbfbad 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -2481,9 +2481,9 @@ std::vector Gather::vjp( std::vector vjps; for (int argnum : argnums) { if (argnum > 0) { - // Grads w.r.t. indices are zero - vjps.push_back( - zeros(primals[argnum].shape(), primals[argnum].dtype(), stream())); + throw std::invalid_argument( + "[gather] Cannot calculate VJP with respect to indices. " + "Use mx.stop_gradient on indices to make this explicit."); } else { auto src = zeros_like(primals[0], stream()); std::vector inds(primals.begin() + 1, primals.end()); @@ -2499,7 +2499,8 @@ std::vector Gather::jvp( const std::vector& argnums) { if (argnums.size() > 1 || argnums[0] != 0) { throw std::invalid_argument( - "[gather] Cannot calculate JVP with respect to indices."); + "[gather] Cannot calculate JVP with respect to indices. " + "Use mx.stop_gradient on indices to make this explicit."); } std::vector inds(primals.begin() + 1, primals.end()); return {gather(tangents[0], inds, axes_, slice_sizes_, stream())}; @@ -2546,9 +2547,9 @@ std::vector GatherAxis::vjp( std::vector vjps; for (int argnum : argnums) { if (argnum > 0) { - // Grads w.r.t. indices are zero - vjps.push_back( - zeros(primals[argnum].shape(), primals[argnum].dtype(), stream())); + throw std::invalid_argument( + "[gather_axis] Cannot calculate VJP with respect to indices. " + "Use mx.stop_gradient on indices to make this explicit."); } else { auto src = zeros_like(primals[0], stream()); vjps.push_back(array( @@ -2567,7 +2568,8 @@ std::vector GatherAxis::jvp( const std::vector& argnums) { if (argnums.size() > 1 || argnums[0] != 0) { throw std::invalid_argument( - "[gather_axis] Cannot calculate JVP with respect to indices."); + "[gather_axis] Cannot calculate JVP with respect to indices. " + "Use mx.stop_gradient on indices to make this explicit."); } return {take_along_axis(tangents[0], primals[1], axis_, stream())}; } @@ -4483,7 +4485,8 @@ std::vector Scatter::vjp( } } else { throw std::invalid_argument( - "[scatter] Cannot calculate VJP with respect to indices."); + "[scatter] Cannot calculate VJP with respect to indices. " + "Use mx.stop_gradient on indices to make this explicit."); } } return vjps; @@ -4595,7 +4598,8 @@ std::vector ScatterAxis::vjp( vjps.push_back(take_along_axis(cotangents[0], indices, axis_, stream())); } else { throw std::invalid_argument( - "[scatter_axis] Cannot calculate VJP with respect to indices."); + "[scatter_axis] Cannot calculate VJP with respect to indices. " + "Use mx.stop_gradient on indices to make this explicit."); } } return vjps; @@ -4608,7 +4612,8 @@ std::vector ScatterAxis::jvp( for (auto arg : argnums) { if (arg == 1) { throw std::invalid_argument( - "[scatter_axis] Cannot calculate JVP with respect to indices."); + "[scatter_axis] Cannot calculate JVP with respect to indices. " + "Use mx.stop_gradient on indices to make this explicit."); } } if (argnums.size() == 2) { @@ -6010,7 +6015,8 @@ std::vector GatherMM::vjp( stream())); } else { throw std::invalid_argument( - "[GatherMM] Cannot calculate VJP with respect to indices."); + "[GatherMM] Cannot calculate VJP with respect to indices. " + "Use mx.stop_gradient on indices to make this explicit."); } } return vjps; diff --git a/python/tests/test_autograd.py b/python/tests/test_autograd.py index fa6aed6f09..1aa3698753 100644 --- a/python/tests/test_autograd.py +++ b/python/tests/test_autograd.py @@ -444,6 +444,83 @@ def fun(x, idx): self.assertTrue(mx.array_equal(dfdx, mx.array([1.0, 1.0]))) self.assertEqual(dfdx.dtype, mx.float32) + def test_index_vjp_requires_stop_gradient(self): + msg = "stop_gradient" + x = mx.array([1.0, 2.0, 3.0, 4.0]) + idx = mx.array([1, 3]) + updates = mx.array([5.0, 6.0]) + x_axis = x[:, None] + idx_axis = idx[:, None] + updates_axis = updates[:, None] + + def gather_fun(x, idx): + return mx.take(x, idx) + + with self.assertRaisesRegex(ValueError, msg): + mx.vjp(gather_fun, [x, idx], [mx.ones((2,))]) + + def gather_axis_fun(x, idx): + return mx.take_along_axis(x, idx, axis=0) + + with self.assertRaisesRegex(ValueError, msg): + mx.vjp(gather_axis_fun, [x_axis, idx_axis], [mx.ones((2, 1))]) + + def scatter_fun(x, idx, updates): + return x.at[idx].add(updates) + + with self.assertRaisesRegex(ValueError, msg): + mx.vjp(scatter_fun, [x, idx, updates], [mx.ones((4,))]) + + def scatter_axis_fun(x, idx, updates): + return mx.put_along_axis(x, idx, updates, axis=0) + + with self.assertRaisesRegex(ValueError, msg): + mx.vjp( + scatter_axis_fun, + [x_axis, idx_axis, updates_axis], + [mx.ones((4, 1))], + ) + + def test_stop_gradient_computed_indices(self): + def gather_fun(w): + idx = mx.stop_gradient(mx.argsort(w)[:2]) + return mx.take(w, idx).sum() + + grad = mx.grad(gather_fun)(mx.array([4.0, 3.0, 2.0, 1.0])) + self.assertTrue(mx.array_equal(grad, mx.array([0.0, 0.0, 1.0, 1.0]))) + + def gather_axis_fun(w): + idx = mx.stop_gradient(mx.argsort(w, axis=1)[:, :1]) + return mx.take_along_axis(w, idx, axis=1).sum() + + grad = mx.grad(gather_axis_fun)( + mx.array([[3.0, 2.0, 1.0], [1.0, 3.0, 2.0]]) + ) + self.assertTrue( + mx.array_equal( + grad, + mx.array([[0.0, 0.0, 1.0], [1.0, 0.0, 0.0]]), + ) + ) + + def scatter_fun(w): + idx = mx.stop_gradient(mx.argsort(w)[:2]) + out = mx.zeros((4,)) + out[idx] = w[:2] + return out.sum() + + grad = mx.grad(scatter_fun)(mx.array([4.0, 3.0, 2.0, 1.0])) + self.assertTrue(mx.array_equal(grad, mx.array([1.0, 1.0, 0.0, 0.0]))) + + def scatter_axis_fun(w): + idx = mx.stop_gradient(mx.argsort(w, axis=1)[:, :1]) + updates = w.sum(axis=1, keepdims=True) + out = mx.put_along_axis(mx.zeros((3, 3)), idx, updates, axis=1) + return out.sum() + + grad = mx.grad(scatter_axis_fun)(mx.ones((3, 3))) + self.assertTrue(mx.array_equal(grad, mx.ones((3, 3)))) + def test_scatter_add_vjp(self): def fun(src, updates): x = src.at[mx.array([1, 3])].add(updates) diff --git a/python/tests/test_blas.py b/python/tests/test_blas.py index 0db6257f93..580846e8e6 100644 --- a/python/tests/test_blas.py +++ b/python/tests/test_blas.py @@ -1239,6 +1239,24 @@ def f_test(a, b): self.assertEqual(r.shape, t.shape) self.assertTrue(mx.allclose(r, t, atol=1e-4).item()) + def test_gather_matmul_index_vjp_requires_stop_gradient(self): + a = mx.ones((4, 1, 2, 2)) + b = mx.ones((4, 1, 2, 2)) + + def fun(w): + indices = mx.reshape(mx.argsort(w)[:2], (1, 2)) + return mx.gather_mm(a, b, indices, indices).sum() + + with self.assertRaisesRegex(ValueError, "stop_gradient"): + mx.grad(fun)(mx.array([3.0, 1.0, 2.0, 0.0])) + + def fun_stopped(w): + indices = mx.stop_gradient(mx.reshape(mx.argsort(w)[:2], (1, 2))) + return mx.gather_mm(a, b, indices, indices).sum() + + grad = mx.grad(fun_stopped)(mx.array([3.0, 1.0, 2.0, 0.0])) + self.assertTrue(mx.array_equal(grad, mx.zeros((4,)))) + def test_gather_mm_sorted(self): def gather_mm_ref(a, b, rhs): b = b[rhs] From c4724e83f8017e4d82cb34b6227dd64f0a2eec2c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fukan=20Veziro=C4=9Flu?= Date: Thu, 9 Jul 2026 04:44:49 +0300 Subject: [PATCH 2/5] Clarify stop_gradient guidance in index autodiff errors --- mlx/primitives.cpp | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index 9c96cbfbad..061c4e5481 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -2483,7 +2483,7 @@ std::vector Gather::vjp( if (argnum > 0) { throw std::invalid_argument( "[gather] Cannot calculate VJP with respect to indices. " - "Use mx.stop_gradient on indices to make this explicit."); + "Use stop_gradient on indices to stop gradients from being computed."); } else { auto src = zeros_like(primals[0], stream()); std::vector inds(primals.begin() + 1, primals.end()); @@ -2500,7 +2500,7 @@ std::vector Gather::jvp( if (argnums.size() > 1 || argnums[0] != 0) { throw std::invalid_argument( "[gather] Cannot calculate JVP with respect to indices. " - "Use mx.stop_gradient on indices to make this explicit."); + "Use stop_gradient on indices to stop gradients from being computed."); } std::vector inds(primals.begin() + 1, primals.end()); return {gather(tangents[0], inds, axes_, slice_sizes_, stream())}; @@ -2549,7 +2549,7 @@ std::vector GatherAxis::vjp( if (argnum > 0) { throw std::invalid_argument( "[gather_axis] Cannot calculate VJP with respect to indices. " - "Use mx.stop_gradient on indices to make this explicit."); + "Use stop_gradient on indices to stop gradients from being computed."); } else { auto src = zeros_like(primals[0], stream()); vjps.push_back(array( @@ -2569,7 +2569,7 @@ std::vector GatherAxis::jvp( if (argnums.size() > 1 || argnums[0] != 0) { throw std::invalid_argument( "[gather_axis] Cannot calculate JVP with respect to indices. " - "Use mx.stop_gradient on indices to make this explicit."); + "Use stop_gradient on indices to stop gradients from being computed."); } return {take_along_axis(tangents[0], primals[1], axis_, stream())}; } @@ -4486,7 +4486,7 @@ std::vector Scatter::vjp( } else { throw std::invalid_argument( "[scatter] Cannot calculate VJP with respect to indices. " - "Use mx.stop_gradient on indices to make this explicit."); + "Use stop_gradient on indices to stop gradients from being computed."); } } return vjps; @@ -4599,7 +4599,7 @@ std::vector ScatterAxis::vjp( } else { throw std::invalid_argument( "[scatter_axis] Cannot calculate VJP with respect to indices. " - "Use mx.stop_gradient on indices to make this explicit."); + "Use stop_gradient on indices to stop gradients from being computed."); } } return vjps; @@ -4613,7 +4613,7 @@ std::vector ScatterAxis::jvp( if (arg == 1) { throw std::invalid_argument( "[scatter_axis] Cannot calculate JVP with respect to indices. " - "Use mx.stop_gradient on indices to make this explicit."); + "Use stop_gradient on indices to stop gradients from being computed."); } } if (argnums.size() == 2) { @@ -6016,7 +6016,7 @@ std::vector GatherMM::vjp( } else { throw std::invalid_argument( "[GatherMM] Cannot calculate VJP with respect to indices. " - "Use mx.stop_gradient on indices to make this explicit."); + "Use stop_gradient on indices to stop gradients from being computed."); } } return vjps; From 279c3a52a92e0c69af2b29594a88324b4b4beceb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fukan=20Veziro=C4=9Flu?= Date: Thu, 9 Jul 2026 04:52:10 +0300 Subject: [PATCH 3/5] Add stop_gradient guidance to mask autodiff errors --- mlx/primitives.cpp | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index 061c4e5481..ae434af94d 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -4721,7 +4721,8 @@ std::vector MaskedScatter::vjp( vjps.push_back(reshape(gsrc_flat, src.shape(), s)); } else { throw std::invalid_argument( - "[masked_scatter] Cannot calculate VJP with respect to mask."); + "[masked_scatter] Cannot calculate VJP with respect to mask. " + "Use stop_gradient on mask to stop gradients from being computed."); } } return vjps; @@ -5753,7 +5754,8 @@ std::vector BlockMaskedMM::vjp( if ((needs_lhs_mask_vjp && primals[op_mask_idx].dtype() == bool_) || (needs_rhs_mask_vjp && primals[op_mask_idx + 1].dtype() == bool_)) { throw std::invalid_argument( - "[BlockMaskedMM] Cannot calculate VJP with respect to boolean masks."); + "[BlockMaskedMM] Cannot calculate VJP with respect to boolean masks. " + "Use stop_gradient on masks to stop gradients from being computed."); } auto expand_mask = [&](array mask, int Y, int X) { @@ -5950,7 +5952,8 @@ std::vector BlockMaskedMM::vjp( } else { throw std::invalid_argument( - "[BlockMaskedMM] Cannot calculate VJP with respect to masks."); + "[BlockMaskedMM] Cannot calculate VJP with respect to masks. " + "Use stop_gradient on masks to stop gradients from being computed."); } } return vjps; From 8b8bb8f5363f68e1d047aaf92a1a2964bf191978 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fukan=20Veziro=C4=9Flu?= Date: Thu, 9 Jul 2026 05:06:03 +0300 Subject: [PATCH 4/5] Apply black formatting to autograd test --- python/tests/test_autograd.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/python/tests/test_autograd.py b/python/tests/test_autograd.py index 1aa3698753..2f72fba113 100644 --- a/python/tests/test_autograd.py +++ b/python/tests/test_autograd.py @@ -493,9 +493,7 @@ def gather_axis_fun(w): idx = mx.stop_gradient(mx.argsort(w, axis=1)[:, :1]) return mx.take_along_axis(w, idx, axis=1).sum() - grad = mx.grad(gather_axis_fun)( - mx.array([[3.0, 2.0, 1.0], [1.0, 3.0, 2.0]]) - ) + grad = mx.grad(gather_axis_fun)(mx.array([[3.0, 2.0, 1.0], [1.0, 3.0, 2.0]])) self.assertTrue( mx.array_equal( grad, From 4715c89dbf64ee942ed3c114fbc5652bd8896b87 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Do=C4=9Fukan=20Veziro=C4=9Flu?= Date: Thu, 9 Jul 2026 05:20:02 +0300 Subject: [PATCH 5/5] Stop gradients through indexing test indices --- python/tests/test_array.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tests/test_array.py b/python/tests/test_array.py index dfab616716..264a354831 100644 --- a/python/tests/test_array.py +++ b/python/tests/test_array.py @@ -1205,7 +1205,7 @@ def test_indexing_grad(self): ind = mx.array([0, 1, 0]).astype(mx.float32) def index_fn(x, ind): - return x[ind.astype(mx.int32)].sum() + return x[mx.stop_gradient(ind.astype(mx.int32))].sum() grad_x, grad_ind = mx.grad(index_fn, argnums=(0, 1))(x, ind) expected = mx.array([[2, 2], [1, 1]])