huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-19 19:58:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add the alloc_uninit function. (#1901)

Browse Source 
* Add the alloc_uninit function.

* Dummy metal fix.

* Lazy initialization.
This commit is contained in:
Laurent Mazare
2024-03-22 07:25:23 +01:00
committed by GitHub
parent a00e24d752
commit 6708870e63
9 changed files with 154 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
candle-core/src/metal_backend.rs
View File

@ -1886,6 +1886,16 @@ impl BackendDevice for MetalDevice {
self.device.registry_id() == rhs.device.registry_id()
}
unsafe fn alloc_uninit(&self, shape: &Shape, dtype: DType) -> Result<MetalStorage> {
let buffer = self.new_buffer(shape.elem_count(), dtype, "alloc-uninit")?;
Ok(MetalStorage::new(
buffer,
self.clone(),
shape.elem_count(),
dtype,
))
}
fn zeros_impl(&self, shape: &Shape, dtype: DType) -> Result<MetalStorage> {
let size = shape.elem_count() * dtype.size_in_bytes();
let buffer = self.allocate_zeros(size)?;