diff options
| author | Laurent Mazare <laurent.mazare@gmail.com> | 2024-03-27 16:30:07 +0100 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2024-03-27 16:30:07 +0100 |
| commit | ab86cd37c8fd944df351d8c7ca0e93376634a332 (patch) | |
| tree | 15f9c8809c1c7c7bd449ab44fefcca24a92fb295 /candle-metal-kernels | |
| parent | a9abde5f930914ef7ef2d504728f742f80468961 (diff) | |
| download | candle-ab86cd37c8fd944df351d8c7ca0e93376634a332.tar.gz candle-ab86cd37c8fd944df351d8c7ca0e93376634a332.tar.bz2 candle-ab86cd37c8fd944df351d8c7ca0e93376634a332.zip | |
Support i64 in index-select on metal. (#1951)
* Support i64 in index-select on metal.
* Add some testing of index-select for all dtypes.
Diffstat (limited to 'candle-metal-kernels')
| -rw-r--r-- | candle-metal-kernels/src/indexing.metal | 8 |
1 files changed, 7 insertions, 1 deletions
diff --git a/candle-metal-kernels/src/indexing.metal b/candle-metal-kernels/src/indexing.metal index ad4a8605..762b42be 100644 --- a/candle-metal-kernels/src/indexing.metal +++ b/candle-metal-kernels/src/indexing.metal @@ -187,6 +187,12 @@ kernel void NAME( \ } +INDEX_OP(is_i64_f32, int64_t, float) +INDEX_OP(is_i64_f16, int64_t, half) +#if defined(__HAVE_BFLOAT__) +INDEX_OP(is_i64_bf16, int64_t, bfloat) +#endif + INDEX_OP(is_u32_f32, uint32_t, float) INDEX_OP(is_u32_f16, uint32_t, half) #if defined(__HAVE_BFLOAT__) @@ -242,4 +248,4 @@ INDEX_ADD_OP(ia_u8_u32, uint8_t, uint32_t) INDEX_ADD_OP(ia_u8_u8, uint8_t, uint8_t) #if defined(__HAVE_BFLOAT__) INDEX_ADD_OP(ia_u8_bf16, uint8_t, bfloat) -#endif
\ No newline at end of file +#endif |