diff options
Diffstat (limited to 'src/device.rs')
| -rw-r--r-- | src/device.rs | 30 |
1 files changed, 2 insertions, 28 deletions
diff --git a/src/device.rs b/src/device.rs index 3677cfff..af538c6c 100644 --- a/src/device.rs +++ b/src/device.rs @@ -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!() } |