Adding matmul?

This commit is contained in:
Nicolas Patry
2023-06-21 16:52:35 +02:00
parent 87a37b3bf3
commit ce977b489e
7 changed files with 243 additions and 3 deletions

View File

@ -25,6 +25,7 @@ pub trait WithDType: Sized + Copy {
}
fn cpu_storage_as_slice(s: &CpuStorage) -> Result<&[Self]>;
fn cpu_storage_as_mut_slice(s: &mut CpuStorage) -> Result<&mut [Self]>;
}
macro_rules! with_dtype {
@ -45,6 +46,16 @@ macro_rules! with_dtype {
}),
}
}
fn cpu_storage_as_mut_slice(s: &mut CpuStorage) -> Result<&mut [Self]> {
match s {
CpuStorage::$dtype(data) => Ok(data),
_ => Err(Error::UnexpectedDType {
expected: DType::$dtype,
got: s.dtype(),
}),
}
}
}
};
}