zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-06 16:51:24 +08:00
Code Issues Actions 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

38
python/src/ops.cpp
View File

@@ -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");
}