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

Move more bits to the backend part.

Browse Source
This commit is contained in:
laurent
2023-06-21 10:34:51 +01:00
parent eb52b9b343
commit 68f525f321
2 changed files with 32 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
src/cpu_backend.rs
View File

@ -68,7 +68,7 @@ impl CpuStorage {
// same if it helps.
// https://github.com/ggerganov/llama.cpp/blob/aacdbd40562684665b6f7b8ba6695b7a2088bbb0/ggml.c#L7895
match (self, rhs) {
(CpuStorage::F32(lhs), CpuStorage::F32(rhs)) => {
(Self::F32(lhs), Self::F32(rhs)) => {
let lhs_index = StridedIndex::new(shape.dims(), lhs_stride);
let rhs_index = StridedIndex::new(shape.dims(), rhs_stride);
let data = lhs_index
@ -77,7 +77,7 @@ impl CpuStorage {
.collect();
Ok(Self::F32(data))
}
(CpuStorage::F64(lhs), CpuStorage::F64(rhs)) => {
(Self::F64(lhs), Self::F64(rhs)) => {
let lhs_index = StridedIndex::new(shape.dims(), lhs_stride);
let rhs_index = StridedIndex::new(shape.dims(), rhs_stride);
let data = lhs_index
@ -96,4 +96,32 @@ impl CpuStorage {
}
}
}
pub(crate) fn ones_impl(shape: &Shape, dtype: DType) -> Self {
let elem_count = shape.elem_count();
match dtype {
DType::F32 => {
let data = vec![1f32; elem_count];
Self::F32(data)
}
DType::F64 => {
let data = vec![1f64; elem_count];
Self::F64(data)
}
}
}
pub(crate) fn zeros_impl(shape: &Shape, dtype: DType) -> Self {
let elem_count = shape.elem_count();
match dtype {
DType::F32 => {
let data = vec![0f32; elem_count];
Self::F32(data)
}
DType::F64 => {
let data = vec![0f64; elem_count];
Self::F64(data)
}
}
}
}

30
src/device.rs
View File

@ -56,20 +56,7 @@ impl<S: crate::WithDType, const N: usize, const M: usize> NdArray for &[[S; N];
impl Device {
pub(crate) fn ones(&self, shape: &Shape, dtype: DType) -> Storage {
match self {
Device::Cpu => {
let elem_count = shape.elem_count();
let storage = match dtype {
DType::F32 => {
let data = vec![1f32; elem_count];
CpuStorage::F32(data)
}
DType::F64 => {
let data = vec![1f64; elem_count];
CpuStorage::F64(data)
}
};
Storage::Cpu(storage)
}
Device::Cpu => Storage::Cpu(CpuStorage::ones_impl(shape, dtype)),
Device::Cuda { gpu_id: _ } => {
todo!()
}
@ -78,20 +65,7 @@ impl Device {
pub(crate) fn zeros(&self, shape: &Shape, dtype: DType) -> Storage {
match self {
Device::Cpu => {
let elem_count = shape.elem_count();
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)
}
Device::Cpu => Storage::Cpu(CpuStorage::zeros_impl(shape, dtype)),
Device::Cuda { gpu_id: _ } => {
todo!()
}