Fix abs jvp for complex inputs

mlx/primitives.cpp

@@ -221,7 +221,9 @@ std::vector<array> Abs::vjp(
     const std::vector<array>& cotangents,
     const std::vector<int>& argnums,
     const std::vector<array>&) {
-  return jvp(primals, cotangents, argnums);
+  assert(primals.size() == 1);
+  assert(argnums.size() == 1);
+  return {multiply(cotangents[0], sign(primals[0], stream()), stream())};
 }
 
 std::vector<array> Abs::jvp(
@@ -230,7 +232,14 @@ std::vector<array> Abs::jvp(
     const std::vector<int>& argnums) {
   assert(primals.size() == 1);
   assert(argnums.size() == 1);
-  return {multiply(tangents[0], sign(primals[0], stream()), stream())};
+  auto s = sign(primals[0], stream());
+  if (issubdtype(primals[0].dtype(), complexfloating)) {
+    // |z| is real-valued, so its directional derivative is real:
+    //   d|z| = Re(conj(z) * t) / |z| = Re(conj(sign(z)) * t).
+    return {real(
+        multiply(tangents[0], conjugate(s, stream()), stream()), stream())};
+  }
+  return {multiply(tangents[0], s, stream())};
 }
 
 std::pair<std::vector<array>, std::vector<int>> Abs::vmap(

python/tests/test_autograd.py

@@ -1168,6 +1168,34 @@ def test_complex_log_vjp(self):
         # Check against hand-computed vjps
         self.assertTrue(mx.allclose(vjps[0], expected))
 
+    def test_complex_abs_grad(self):
+        mx.random.seed(0)
+        primal = mx.random.normal((3, 4, 5), dtype=mx.complex64)
+        # guard against values too close to the origin where |z| is not smooth
+        primal = mx.where(abs(primal) < 1e-3, 1e-3 + 0j, primal)
+
+        # |z| is real-valued, so its jvp is real:
+        #   d|z| = Re(conj(z) * t) / |z|
+        tangent = mx.random.normal(primal.shape, dtype=mx.complex64)
+        _, (jvp,) = mx.jvp(mx.abs, [primal], [tangent])
+        expected = mx.real(mx.conj(primal) * tangent) / mx.abs(primal)
+        self.assertEqual(jvp.dtype, mx.float32)
+        self.assertTrue(mx.allclose(jvp, expected, atol=1e-5))
+
+        # The vjp's real and imaginary parts are the gradients w.r.t. Re(z) and
+        # Im(z); for a real cotangent this is cotangent * sign(z).
+        cotangent = mx.random.normal(primal.shape)
+        _, (vjp,) = mx.vjp(mx.abs, [primal], [cotangent])
+        self.assertTrue(
+            mx.allclose(vjp, cotangent * (primal / mx.abs(primal)), atol=1e-5)
+        )
+
+        # Real inputs are unaffected.
+        x = mx.random.normal((10,))
+        t = mx.random.normal((10,))
+        _, (jvp,) = mx.jvp(mx.abs, [x], [t])
+        self.assertTrue(mx.allclose(jvp, mx.sign(x) * t))
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()