zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-19 00:04:41 +08:00
Code Issues Actions 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

19
python/src/ops.cpp
View File

@@ -930,6 +930,25 @@ void init_ops(nb::module_& m) {
Returns:
array: The inverse tangent of ``a``.
)pbdoc");
m.def(
"arctan2",
&mlx::core::arctan2,
nb::arg(),
nb::arg(),
nb::kw_only(),
"stream"_a = nb::none(),
nb::sig(
"def arctan2(a: array, b: array, /, *, stream: Union[None, Stream, Device] = None) -> array"),
R"pbdoc(
Element-wise inverse tangent of the ratio of two arrays.
Args:
a (array): Input array.
b (array): Input array.
Returns:
array: The inverse tangent of the ratio of ``a`` and ``b``.
)pbdoc");
m.def(
"sinh",
&mlx::core::sinh,

1
python/tests/test_ops.py
View File

@@ -1273,6 +1273,7 @@ class TestOps(mlx_tests.MLXTestCase):
"arcsin": lambda primal, cotan: cotan / np.sqrt(1.0 - primal**2),
"arccos": lambda primal, cotan: -cotan / np.sqrt(1.0 - primal**2),
"arctan": lambda primal, cotan: cotan / (1.0 + primal**2),
"arctan2": lambda primal, cotan: cotan / (1.0 + primal**2),
"arcsinh": lambda primal, cotan: cotan / np.sqrt(primal**2 + 1),
"arccosh": lambda primal, cotan: cotan / np.sqrt(primal**2 - 1),
"arctanh": lambda primal, cotan: cotan / (1.0 - primal**2),