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

Only check nans on non-integral types in simd_reduce_impl.

Browse Source
This commit is contained in:
Joona Havukainen
2025-07-07 18:24:30 -07:00
parent 9794ec6b8e
commit c7af3016eb
1 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
mlx/backend/metal/kernels/reduction/ops.h
View File

@@ -186,7 +186,12 @@ struct Max {
DEFINE_SIMD_REDUCE() DEFINE_SIMD_REDUCE()
template <typename T> template <typename T>
T simd_reduce_impl(T val) { metal::enable_if_t<metal::is_integral_v<T>, T> simd_reduce_impl(T val) {
return simd_max(val);
}
template <typename T>
metal::enable_if_t<!metal::is_integral_v<T>, T> simd_reduce_impl(T val) {
if (simd_any(val != val)) { if (simd_any(val != val)) {
return static_cast<T>(NAN); return static_cast<T>(NAN);
} }