zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 01:17:26 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

real and imag properties (#2189)

Browse Source
This commit is contained in:
Awni Hannun 2025-05-15 18:17:50 -07:00 committed by GitHub
parent c1eb9d05d9
commit a2cadb8218
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 23 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
docs/src/python/array.rst
View File

@ -19,6 +19,8 @@ Array
array.ndim
array.shape
array.size
array.real
array.imag
array.abs
array.all
array.any

12
python/src/array.cpp
View File

@ -319,6 +319,18 @@ void init_array(nb::module_& m) {
R"pbdoc(
The array's :class:`Dtype`.
)pbdoc")
.def_prop_ro(
"real",
[](const mx::array& a) { return mx::real(a); },
R"pbdoc(
The real part of a complex array.
)pbdoc")
.def_prop_ro(
"imag",
[](const mx::array& a) { return mx::imag(a); },
R"pbdoc(
The imaginary part of a complex array.
)pbdoc")
.def(
"item",
&to_scalar,

9
python/tests/test_array.py
View File

@ -2022,6 +2022,15 @@ class TestArray(mlx_tests.MLXTestCase):
with self.assertRaises(ValueError):
mx.add(y, x)
def test_real_imag(self):
x = mx.array([1.0])
self.assertEqual(x.real.item(), 1.0)
self.assertEqual(x.imag.item(), 0.0)
x = mx.array([1.0 + 1.0j])
self.assertEqual(x.imag.item(), 1.0)
self.assertEqual(x.real.item(), 1.0)
if __name__ == "__main__":
unittest.main()