huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-15 02:16:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add a few metal gather ops. (#2740)

Browse Source 
* Add a few metal gather ops.

* Fix some compilation issues.

* Adjust the tolerance.
This commit is contained in:
Laurent Mazare
2025-01-25 23:31:03 +01:00
committed by GitHub
parent 333d94a19a
commit 1a32107fab
5 changed files with 17 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
candle-core/src/metal_backend/mod.rs
View File

@ -1245,6 +1245,12 @@ impl BackendStorage for MetalStorage {
(DType::U32, DType::F16) => "gather_u32_f16",
(DType::U32, DType::BF16) => "gather_u32_bf16",
(DType::U32, DType::U32) => "gather_u32_u32",
(DType::U32, DType::I64) => "gather_u32_i64",
(DType::I64, DType::F32) => "gather_i64_f32",
(DType::I64, DType::F16) => "gather_i64_f16",
(DType::I64, DType::BF16) => "gather_i64_bf16",
(DType::I64, DType::U32) => "gather_i64_u32",
(DType::I64, DType::I64) => "gather_i64_i64",
(left, right) => crate::bail!("Metal gather {left:?} {right:?} not implemented"),
};
let command_buffer = self.device.command_buffer()?;