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

Abstract the gradient storage.

Browse Source
This commit is contained in:
laurent
2023-06-21 14:29:48 +01:00
parent 68f525f321
commit 7adffafeda
5 changed files with 87 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
src/device.rs
View File

@ -54,27 +54,36 @@ 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 {
pub(crate) fn ones(&self, shape: &Shape, dtype: DType) -> Result<Storage> {
match self {
Device::Cpu => Storage::Cpu(CpuStorage::ones_impl(shape, dtype)),
Device::Cpu => {
let storage = Storage::Cpu(CpuStorage::ones_impl(shape, dtype));
Ok(storage)
}
Device::Cuda { gpu_id: _ } => {
todo!()
}
}
}
pub(crate) fn zeros(&self, shape: &Shape, dtype: DType) -> Storage {
pub(crate) fn zeros(&self, shape: &Shape, dtype: DType) -> Result<Storage> {
match self {
Device::Cpu => Storage::Cpu(CpuStorage::zeros_impl(shape, dtype)),
Device::Cpu => {
let storage = Storage::Cpu(CpuStorage::zeros_impl(shape, dtype));
Ok(storage)
}
Device::Cuda { gpu_id: _ } => {
todo!()
}
}
}
pub(crate) fn tensor<A: NdArray>(&self, array: A) -> Storage {
pub(crate) fn tensor<A: NdArray>(&self, array: A) -> Result<Storage> {
match self {
Device::Cpu => Storage::Cpu(array.to_cpu_storage()),
Device::Cpu => {
let storage = Storage::Cpu(array.to_cpu_storage());
Ok(storage)
}
Device::Cuda { gpu_id: _ } => {
todo!()
}