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

fix isinf for integer types (#494)

Browse Source
This commit is contained in:
Awni Hannun
2024-01-19 05:31:10 -08:00
committed by GitHub
parent 550d4bf7c0
commit c4ec836523
2 changed files with 36 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
mlx/ops.cpp
View File

@@ -1105,18 +1105,30 @@ array array_equal(
}
array isnan(const array& a, StreamOrDevice s /* = {} */) {
if (is_integral(a.dtype())) {
return full(a.shape(), false, bool_, s);
}
return not_equal(a, a, s);
}
array isinf(const array& a, StreamOrDevice s /* = {} */) {
if (is_integral(a.dtype())) {
return full(a.shape(), false, bool_, s);
}
return equal(a, array(std::numeric_limits<float>::infinity(), a.dtype()), s);
}
array isposinf(const array& a, StreamOrDevice s) {
if (is_integral(a.dtype())) {
return full(a.shape(), false, bool_, s);
}
return equal(a, array(std::numeric_limits<float>::infinity(), a.dtype()), s);
}
array isneginf(const array& a, StreamOrDevice s) {
if (is_integral(a.dtype())) {
return full(a.shape(), false, bool_, s);
}
return equal(a, array(-std::numeric_limits<float>::infinity(), a.dtype()), s);
}