diff options
| author | Thomas Santerre <thomas@santerre.xyz> | 2024-04-04 12:39:06 -0400 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2024-04-04 18:39:06 +0200 |
| commit | 5aebe53dd2470db731bd9ce2d65914f86f1542f7 (patch) | |
| tree | 0046d3cde6a2a521e4122aacc8ae5bdce176f399 /candle-metal-kernels/src/reduce.metal | |
| parent | f76bb7794aa8659c5023797979a3392cdfc01f32 (diff) | |
| download | candle-5aebe53dd2470db731bd9ce2d65914f86f1542f7.tar.gz candle-5aebe53dd2470db731bd9ce2d65914f86f1542f7.tar.bz2 candle-5aebe53dd2470db731bd9ce2d65914f86f1542f7.zip | |
update dtypes checks for several metal operations (#2010)
Diffstat (limited to 'candle-metal-kernels/src/reduce.metal')
| -rw-r--r-- | candle-metal-kernels/src/reduce.metal | 4 |
1 files changed, 4 insertions, 0 deletions
diff --git a/candle-metal-kernels/src/reduce.metal b/candle-metal-kernels/src/reduce.metal index 561d1744..acb69299 100644 --- a/candle-metal-kernels/src/reduce.metal +++ b/candle-metal-kernels/src/reduce.metal @@ -484,9 +484,13 @@ ARGMAX(fast_argmax_i64_strided, int64_t, INT_MIN) #if defined(__HAVE_BFLOAT__) REDUCE(x + y, fast_sum_bf16, bfloat, 0) +REDUCE(x + y, fast_sum_bf16_strided, half, 0) REDUCE(x * y, fast_mul_bf16, bfloat, 1) +REDUCE(x * y, fast_mul_bf16_strided, bfloat, 1) REDUCE(MAX(x, y), fast_max_bf16, bfloat, -HUGE_VALBF) +REDUCE(MAX(x, y), fast_max_bf16_strided, bfloat, -HUGE_VALBF) REDUCE(MIN(x, y), fast_min_bf16, bfloat, HUGE_VALBF) +REDUCE(MIN(x, y), fast_min_bf16_strided, bfloat, HUGE_VALBF) ARGMIN(fast_argmin_bf16, bfloat, HUGE_VALBF) ARGMAX(fast_argmax_bf16, bfloat, -HUGE_VALBF) SOFTMAX(softmax_bf16, bfloat) |