Make index autodiff errors explicitly recommend stop_gradient

mlx/primitives.cpp

@@ -2481,9 +2481,9 @@ std::vector<array> Gather::vjp(
   std::vector<array> 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 stop_gradient on indices to stop gradients from being computed.");
     } else {
       auto src = zeros_like(primals[0], stream());
       std::vector<array> inds(primals.begin() + 1, primals.end());
@@ -2499,7 +2499,8 @@ std::vector<array> Gather::jvp(
     const std::vector<int>& 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 stop_gradient on indices to stop gradients from being computed.");
   }
   std::vector<array> inds(primals.begin() + 1, primals.end());
   return {gather(tangents[0], inds, axes_, slice_sizes_, stream())};
@@ -2546,9 +2547,9 @@ std::vector<array> GatherAxis::vjp(
   std::vector<array> 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 stop_gradient on indices to stop gradients from being computed.");
     } else {
       auto src = zeros_like(primals[0], stream());
       vjps.push_back(array(
@@ -2567,7 +2568,8 @@ std::vector<array> GatherAxis::jvp(
     const std::vector<int>& 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 stop_gradient on indices to stop gradients from being computed.");
   }
   return {take_along_axis(tangents[0], primals[1], axis_, stream())};
 }
@@ -4483,7 +4485,8 @@ std::vector<array> Scatter::vjp(
       }
     } else {
       throw std::invalid_argument(
-          "[scatter] Cannot calculate VJP with respect to indices.");
+          "[scatter] Cannot calculate VJP with respect to indices. "
+          "Use stop_gradient on indices to stop gradients from being computed.");
     }
   }
   return vjps;
@@ -4595,7 +4598,8 @@ std::vector<array> 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 stop_gradient on indices to stop gradients from being computed.");
     }
   }
   return vjps;
@@ -4608,7 +4612,8 @@ std::vector<array> 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 stop_gradient on indices to stop gradients from being computed.");
     }
   }
   if (argnums.size() == 2) {
@@ -4716,7 +4721,8 @@ std::vector<array> 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;
@@ -5748,7 +5754,8 @@ std::vector<array> 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) {
@@ -5945,7 +5952,8 @@ std::vector<array> 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;
@@ -6010,7 +6018,8 @@ std::vector<array> 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 stop_gradient on indices to stop gradients from being computed.");
     }
   }
   return vjps;

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]])

python/tests/test_autograd.py

@@ -444,6 +444,81 @@ 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)

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]