From d6d9b248018a461dffaa2b988ab325556e48bf82 Mon Sep 17 00:00:00 2001
From: Cameron Churchwell <53796107+CameronChurchwell@users.noreply.github.com>
Date: Thu, 9 Apr 2026 17:04:46 -0500
Subject: [PATCH] Conjugate VJP and JVP support (#3386)

---
 mlx/primitives.cpp       | 15 +++++++++++++++
 mlx/primitives.h         |  1 +
 python/tests/test_ops.py | 10 ++++++++++
 3 files changed, 26 insertions(+)

diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp
index a861dba4..8e209eeb 100644
--- a/mlx/primitives.cpp
+++ b/mlx/primitives.cpp
@@ -1178,6 +1178,21 @@ std::vector<Shape> Concatenate::output_shapes(
   return {std::move(shape)};
 }
 
+std::vector<array> Conjugate::vjp(
+    const std::vector<array>&,
+    const std::vector<array>& cotangents,
+    const std::vector<int>&,
+    const std::vector<array>&) {
+  return {conjugate(cotangents[0], stream())};
+}
+
+std::vector<array> Conjugate::jvp(
+    const std::vector<array>&,
+    const std::vector<array>& tangents,
+    const std::vector<int>&) {
+  return {conjugate(tangents[0], stream())};
+}
+
 std::pair<std::vector<array>, std::vector<int>> Conjugate::vmap(
     const std::vector<array>& inputs,
     const std::vector<int>& axes) {
diff --git a/mlx/primitives.h b/mlx/primitives.h
index ed580d6a..75fb978d 100644
--- a/mlx/primitives.h
+++ b/mlx/primitives.h
@@ -698,6 +698,7 @@ class Conjugate : public UnaryPrimitive {
   void eval_gpu(const std::vector<array>& inputs, array& out) override;
 
   DEFINE_VMAP()
+  DEFINE_GRADS()
   DEFINE_NAME(Conjugate)
   DEFINE_DEFAULT_IS_EQUIVALENT()
   DEFINE_INPUT_OUTPUT_SHAPE()
diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py
index 12b5d3f3..ccc26bdc 100644
--- a/python/tests/test_ops.py
+++ b/python/tests/test_ops.py
@@ -2888,6 +2888,16 @@ class TestOps(mlx_tests.MLXTestCase):
         out_np = a.conj()
         self.assertTrue(np.array_equal(np.array(out_mlx), out_np))
 
+        b = np.random.normal(size=shape) + 1j * np.random.normal(size=shape)
+        b = b.astype(np.complex64)
+
+        _, vjps = mx.vjp(mx.conj, [mx.array(a)], [mx.array(b)])
+        self.assertTrue(np.array_equal(np.array(vjps[0]), b.conj()))
+
+        out_mlx, jvps = mx.jvp(mx.conj, [mx.array(a)], [mx.array(b)])
+        self.assertTrue(np.array_equal(np.array(out_mlx[0]), a.conj()))
+        self.assertTrue(np.array_equal(np.array(jvps[0]), b.conj()))
+
     def test_view(self):
         # Check scalar
         out = mx.array(1, mx.int8).view(mx.uint8).item()