diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index 6ac88eef6c..ae434af94d 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 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()); @@ -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 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())}; @@ -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 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 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 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 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 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 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 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 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 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 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; 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]]) diff --git a/python/tests/test_autograd.py b/python/tests/test_autograd.py index fa6aed6f09..2f72fba113 100644 --- a/python/tests/test_autograd.py +++ b/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) 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]