zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-05 11:28:12 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Real and Imag (#1490)

Browse Source 
* real and imag

* fix

* fix
This commit is contained in:
Awni Hannun
2024-10-15 16:23:15 -07:00
committed by GitHub
parent 2b8ace6a03
commit 3f86399922
21 changed files with 275 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
python/tests/test_ops.py
View File

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