Fix mx.prod vjp for complex types (#3433)

Author: CameronChurchwell

mlx/primitives.cpp

@@ -3951,7 +3951,7 @@ std::vector<array> 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

python/tests/test_autograd.py

@@ -966,6 +966,27 @@ def test_reduce_jvp(self):
         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()