zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-18 10:26:56 +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

4
mlx/ops.cpp
View File

@@ -1084,6 +1084,10 @@ array isnan(const array& a, StreamOrDevice s /* = {} */) {
return not_equal(a, a, s);
}
array isinf(const array& a, StreamOrDevice s /* = {} */) {
return equal(a, array(std::numeric_limits<float>::infinity(), a.dtype()), s);
}
array where(
const array& condition,
const array& x,

2
mlx/ops.h
View File

@@ -378,6 +378,8 @@ array_equal(const array& a, const array& b, StreamOrDevice s = {}) {
array isnan(const array& a, StreamOrDevice s = {});
array isinf(const array& a, StreamOrDevice s = {});
/** Select from x or y depending on condition. */
array where(
const array& condition,