zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-19 08:38:09 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

implemented isposinf and isneginf in one PR (#470)

Browse Source 
* ran precommit

* updated docs
This commit is contained in:
Yashraj Singh
2024-01-16 20:18:07 +05:30
committed by GitHub
parent a2ffea683a
commit e72458a3fa
6 changed files with 128 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
python/src/ops.cpp
View File

@@ -1856,6 +1856,44 @@ void init_ops(py::module_& m) {
Returns:
array: The boolean array indicating which elements are +/- infinity.
)pbdoc");
m.def(
"isposinf",
&isposinf,
"a"_a,
py::pos_only(),
py::kw_only(),
"stream"_a = none,
R"pbdoc(
isposinf(a: array, stream: Union[None, Stream, Device] = None) -> array
Return a boolean array indicating which elements are positive infinity.
Args:
a (array): Input array.
stream (Union[None, Stream, Device]): Optional stream or device.
Returns:
array: The boolean array indicating which elements are positive infinity.
)pbdoc");
m.def(
"isneginf",
&isneginf,
"a"_a,
py::pos_only(),
py::kw_only(),
"stream"_a = none,
R"pbdoc(
isneginf(a: array, stream: Union[None, Stream, Device] = None) -> array
Return a boolean array indicating which elements are negative infinity.
Args:
a (array): Input array.
stream (Union[None, Stream, Device]): Optional stream or device.
Returns:
array: The boolean array indicating which elements are negative infinity.
)pbdoc");
m.def(
"moveaxis",
&moveaxis,