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

Added ArcTan2 operation (#1079)

Browse Source 
* Added ArcTan2 operation

* Cleanup, bug fixes from code review

* Minor cleanup, fixed Linux tests
This commit is contained in:
Rahul Yedida
2024-05-08 11:35:15 -04:00
committed by GitHub
parent fe96ceee66
commit cc05a281c4
16 changed files with 143 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
mlx/ops.cpp
View File

@@ -2152,6 +2152,14 @@ array arctan(const array& a, StreamOrDevice s /* = {} */) {
a.shape(), dtype, std::make_shared<ArcTan>(to_stream(s)), {input});
}
array arctan2(const array& a, const array& b, StreamOrDevice s /* = {} */) {
auto dtype = at_least_float(promote_types(a.dtype(), b.dtype()));
auto inputs = broadcast_arrays(astype(a, dtype, s), astype(b, dtype, s), s);
auto& shape = inputs[0].shape();
return array(
shape, dtype, std::make_shared<ArcTan2>(to_stream(s)), std::move(inputs));
}
array sinh(const array& a, StreamOrDevice s /* = {} */) {
auto dtype = at_least_float(a.dtype());
auto input = astype(a, dtype, s);