Real and Imag (#1490)

* real and imag

* fix

* fix

python/src/ops.cpp

@@ -4842,4 +4842,42 @@ void init_ops(nb::module_& m) {
           axis (int or tuple(int), optional): The axis or axes along which to
             roll the elements.
       )pbdoc");
+  m.def(
+      "real",
+      [](const ScalarOrArray& a, StreamOrDevice s) {
+        return mlx::core::real(to_array(a), s);
+      },
+      nb::arg(),
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def real(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Returns the real part of a complex array.
+
+        Args:
+            a (array): Input array.
+
+        Returns:
+            array: The real part of ``a``.
+      )pbdoc");
+  m.def(
+      "imag",
+      [](const ScalarOrArray& a, StreamOrDevice s) {
+        return mlx::core::imag(to_array(a), s);
+      },
+      nb::arg(),
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def imag(a: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Returns the imaginary part of a complex array.
+
+        Args:
+            a (array): Input array.
+
+        Returns:
+            array: The imaginary part of ``a``.
+      )pbdoc");
 }

python/tests/test_autograd.py

@@ -590,6 +590,21 @@ class TestAutograd(mlx_tests.MLXTestCase):
         self.assertTrue(mx.allclose(out1[0], out2[0]))
         self.assertTrue(mx.allclose(dout1[0] + 1, dout2[0]))
 
+    def test_complex_vjps(self):
+        def fun(x):
+            return (2.0 * mx.real(x)).sum()
+
+        x = mx.array([0.0 + 1j, 1.0 + 0.0j, 0.5 + 0.5j])
+        dfdx = mx.grad(fun)(x)
+        self.assertTrue(mx.allclose(dfdx, 2 * mx.ones_like(x)))
+
+        def fun(x):
+            return (2.0 * mx.imag(x)).sum()
+
+        x = mx.array([0.0 + 1j, 1.0 + 0.0j, 0.5 + 0.5j])
+        dfdx = mx.grad(fun)(x)
+        self.assertTrue(mx.allclose(dfdx, -2j * mx.ones_like(x)))
+
 
 if __name__ == "__main__":
     unittest.main()

python/tests/test_ops.py

@@ -2680,6 +2680,21 @@ class TestOps(mlx_tests.MLXTestCase):
             y2 = mx.roll(x, s, a)
             self.assertTrue(mx.array_equal(y1, y2).item())
 
+    def test_real_imag(self):
+        x = mx.random.uniform(shape=(4, 4))
+        out = mx.real(x)
+        self.assertTrue(mx.array_equal(x, out))
+
+        out = mx.imag(x)
+        self.assertTrue(mx.array_equal(mx.zeros_like(x), out))
+
+        y = mx.random.uniform(shape=(4, 4))
+        z = x + 1j * y
+        self.assertEqual(mx.real(z).dtype, mx.float32)
+        self.assertTrue(mx.array_equal(mx.real(z), x))
+        self.assertEqual(mx.imag(z).dtype, mx.float32)
+        self.assertTrue(mx.array_equal(mx.imag(z), y))
+
 
 if __name__ == "__main__":
     unittest.main()