zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +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

14
mlx/ops.cpp
View File

@@ -4657,4 +4657,18 @@ array roll(
return roll(a, std::vector<int>{total_shift}, std::vector<int>{axis}, s);
}
array real(const array& a, StreamOrDevice s /* = {} */) {
if (!issubdtype(a.dtype(), complexfloating)) {
return a;
}
return array(a.shape(), float32, std::make_shared<Real>(to_stream(s)), {a});
}
array imag(const array& a, StreamOrDevice s /* = {} */) {
if (!issubdtype(a.dtype(), complexfloating)) {
return zeros_like(a);
}
return array(a.shape(), float32, std::make_shared<Imag>(to_stream(s)), {a});
}
} // namespace mlx::core