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

Flesh out some ops bits.

Browse Source
This commit is contained in:
laurent
2023-06-19 19:28:33 +01:00
parent ce718bb807
commit 8e2c534d1f
3 changed files with 35 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
src/device.rs
View File

@ -1,4 +1,7 @@
use crate::{storage::Storage, DType};
use crate::{
storage::{CpuStorage, Storage},
DType,
};
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum Device {
@ -10,8 +13,17 @@ impl Device {
match self {
Device::Cpu => {
let elem_count: usize = shape.iter().product();
let buffer = vec![0; elem_count * dtype.size_in_bytes()];
Storage::Cpu { dtype, buffer }
let storage = match dtype {
DType::F32 => {
let data = vec![0f32; elem_count];
CpuStorage::F32(data)
}
DType::F64 => {
let data = vec![0f64; elem_count];
CpuStorage::F64(data)
}
};
Storage::Cpu(storage)
}
}
}