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

Fix out-of-bounds default value in logsumexp/softmax (#2213)

Browse Source
This commit is contained in:
Cheng
2025-05-21 23:25:16 +09:00
committed by GitHub
parent 7774b87cbd
commit 79071bfba4
3 changed files with 7 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
mlx/backend/metal/kernels/logsumexp.h
View File

@@ -103,8 +103,8 @@ template <typename T, typename AccT = float, int N_READS = 4>
}
} else {
for (int i = 0; i < N_READS; i++) {
vals[i] = (offset + i < axis_size) ? AccT(in[offset + i])
: Limits<AccT>::finite_min;
vals[i] =
(offset + i < axis_size) ? AccT(in[offset + i]) : Limits<AccT>::min;
}
}
prevmax = maxval;

4
mlx/backend/metal/kernels/softmax.h
View File

@@ -128,8 +128,8 @@ template <typename T, typename AccT = T, int N_READS = SOFTMAX_N_READS>
}
} else {
for (int i = 0; i < N_READS; i++) {
vals[i] = (offset + i < axis_size) ? AccT(in[offset + i])
: Limits<AccT>::finite_min;
vals[i] =
(offset + i < axis_size) ? AccT(in[offset + i]) : Limits<AccT>::min;
}
}
prevmax = maxval;