From 68cf2fddd8de5edd8ab3d926391772b2e2cedad8 Mon Sep 17 00:00:00 2001 From: Cameron Churchwell <53796107+CameronChurchwell@users.noreply.github.com> Date: Tue, 21 Apr 2026 19:35:20 -0500 Subject: [PATCH] Fix mx.prod vjp for complex types (#3433) --- mlx/primitives.cpp | 2 +- python/tests/test_autograd.py | 21 +++++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index 8e209eeb..f3acec57 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -3951,7 +3951,7 @@ std::vector Reduce::vjp( auto p1 = cumprod(x, axis, /*reverse=*/false, /*inclusive=*/false, s); auto p2 = cumprod(x, axis, /*reverse=*/true, /*inclusive=*/false, s); auto exclusive_prod = multiply(p1, p2, s); - return multiply(exclusive_prod, cotan, s); + return multiply(conjugate(exclusive_prod, s), cotan, s); }; // To compute a numerically stable gradient for prod we need an exclusive diff --git a/python/tests/test_autograd.py b/python/tests/test_autograd.py index f0dcc6e1..31ceda1a 100644 --- a/python/tests/test_autograd.py +++ b/python/tests/test_autograd.py @@ -966,6 +966,27 @@ class TestAutograd(mlx_tests.MLXTestCase): out, jout = mx.jvp(mx.max, primals=(a,), tangents=(b,)) self.assertEqual(jout[0].item(), 0) + def test_complex_prod_vjp(self): + def prod(x): + return x.prod(axis=0) + + primal = mx.random.normal((2, 20), dtype=mx.complex64) + cotangent = mx.random.normal((20,), dtype=mx.complex64) + + _, vjps = mx.vjp(prod, [primal], [cotangent]) + + expected = mx.stack( + [mx.conj(primal[1]) * cotangent, mx.conj(primal[0]) * cotangent] + ) + + # Check against hand-computed vjps + self.assertTrue(mx.array_equal(vjps[0], expected)) + + # Ensure that prod agrees with multiply for complex values + _, vjps_multiply = mx.vjp(mx.multiply, [primal[0], primal[1]], [cotangent]) + + self.assertTrue(mx.array_equal(mx.stack(vjps_multiply), vjps[0])) + if __name__ == "__main__": mlx_tests.MLXTestRunner()