Real and Imag (#1490)

* real and imag

* fix

* fix

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()