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

Start adding some ops.

Browse Source
This commit is contained in:
laurent
2023-06-20 08:41:19 +01:00
parent ef6760117f
commit 7a31ba93e4
5 changed files with 83 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
src/dtype.rs
View File

@ -1,4 +1,4 @@
use crate::CpuStorage;
use crate::{CpuStorage, Error, Result};
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum DType {
@ -19,6 +19,8 @@ pub trait WithDType: Sized + Copy {
const DTYPE: DType;
fn to_cpu_storage(data: &[Self]) -> CpuStorage;
fn cpu_storage_as_slice(s: &CpuStorage) -> Result<&[Self]>;
}
impl WithDType for f32 {
@ -27,6 +29,16 @@ impl WithDType for f32 {
fn to_cpu_storage(data: &[Self]) -> CpuStorage {
CpuStorage::F32(data.to_vec())
}
fn cpu_storage_as_slice(s: &CpuStorage) -> Result<&[Self]> {
match s {
CpuStorage::F32(data) => Ok(data),
_ => Err(Error::UnexpectedDType {
expected: DType::F32,
got: s.dtype(),
}),
}
}
}
impl WithDType for f64 {
@ -35,4 +47,14 @@ impl WithDType for f64 {
fn to_cpu_storage(data: &[Self]) -> CpuStorage {
CpuStorage::F64(data.to_vec())
}
fn cpu_storage_as_slice(s: &CpuStorage) -> Result<&[Self]> {
match s {
CpuStorage::F64(data) => Ok(data),
_ => Err(Error::UnexpectedDType {
expected: DType::F64,
got: s.dtype(),
}),
}
}
}