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

Adds isinf (#445)

Browse Source 
* adds isinf

Signed-off-by: matthewfernst <matthew.f.ernst@gmail.com>

* use stream + nits

* typo

---------

Signed-off-by: matthewfernst <matthew.f.ernst@gmail.com>
Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
Matthew Ernst
2024-01-15 19:50:44 -08:00
committed by GitHub
parent 6022d4129e
commit 92a2fdd577
7 changed files with 67 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/src/array.cpp
View File

@@ -1134,12 +1134,6 @@ void init_array(py::module_& m) {
py::kw_only(),
"stream"_a = none,
"See :func:`any`.")
.def(
"isnan",
&mlx::core::isnan,
py::kw_only(),
"stream"_a = none,
"See :func:`isnan`.")
.def(
"moveaxis",
&moveaxis,

20
python/src/ops.cpp
View File

@@ -1836,7 +1836,25 @@ void init_ops(py::module_& m) {
a (array): Input array.
Returns:
array: The array with boolean values indicating which elements are NaN.
array: The boolean array indicating which elements are NaN.
)pbdoc");
m.def(
"isinf",
&mlx::core::isinf,
"a"_a,
py::pos_only(),
py::kw_only(),
"stream"_a = none,
R"pbdoc(
isinf(a: array, stream: Union[None, Stream, Device] = None) -> array
Return a boolean array indicating which elements are +/- inifnity.
Args:
a (array): Input array.
Returns:
array: The boolean array indicating which elements are +/- infinity.
)pbdoc");
m.def(
"moveaxis",