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

Get the addition/multiplication to work.

Browse Source
This commit is contained in:
laurent
2023-06-20 11:07:59 +01:00
parent 6c5fc767a8
commit f5b0aa815a
2 changed files with 62 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
src/storage.rs
View File

@ -119,8 +119,8 @@ impl Storage {
&self,
rhs: &Self,
shape: &Shape,
_lhs_stride: &[usize],
_rhs_stride: &[usize],
lhs_stride: &[usize],
rhs_stride: &[usize],
) -> Result<Self> {
self.same_device(rhs, "add")?;
self.same_dtype(rhs, "add")?;
@ -130,16 +130,22 @@ impl Storage {
// https://github.com/ggerganov/llama.cpp/blob/aacdbd40562684665b6f7b8ba6695b7a2088bbb0/ggml.c#L7895
match (self, rhs) {
(Storage::Cpu(lhs), Storage::Cpu(rhs)) => match (lhs, rhs) {
(CpuStorage::F32(_), CpuStorage::F32(_)) => {
let elem_count = shape.elem_count();
let data = vec![0f32; elem_count];
// TODO: properly fill data with the sum
(CpuStorage::F32(lhs), CpuStorage::F32(rhs)) => {
let lhs_index = StridedIndex::new(shape.dims(), lhs_stride);
let rhs_index = StridedIndex::new(shape.dims(), rhs_stride);
let data = lhs_index
.zip(rhs_index)
.map(|(lhs_i, rhs_i)| lhs[lhs_i] + rhs[rhs_i])
.collect();
Ok(Storage::Cpu(CpuStorage::F32(data)))
}
(CpuStorage::F64(_), CpuStorage::F64(_)) => {
let elem_count = shape.elem_count();
let data = vec![0f64; elem_count];
// TODO: properly fill data with the sum
(CpuStorage::F64(lhs), CpuStorage::F64(rhs)) => {
let lhs_index = StridedIndex::new(shape.dims(), lhs_stride);
let rhs_index = StridedIndex::new(shape.dims(), rhs_stride);
let data = lhs_index
.zip(rhs_index)
.map(|(lhs_i, rhs_i)| lhs[lhs_i] + rhs[rhs_i])
.collect();
Ok(Storage::Cpu(CpuStorage::F64(data)))
}
_ => {
@ -158,12 +164,42 @@ impl Storage {
pub(crate) fn mul_impl(
&self,
rhs: &Self,
_shape: &Shape,
_lhs_stride: &[usize],
_rhs_stride: &[usize],
shape: &Shape,
lhs_stride: &[usize],
rhs_stride: &[usize],
) -> Result<Self> {
self.same_device(rhs, "mul")?;
self.same_dtype(rhs, "mul")?;
todo!()
// TODO: share this code with the add implementation, using a macro or a trait?
match (self, rhs) {
(Storage::Cpu(lhs), Storage::Cpu(rhs)) => match (lhs, rhs) {
(CpuStorage::F32(lhs), CpuStorage::F32(rhs)) => {
let lhs_index = StridedIndex::new(shape.dims(), lhs_stride);
let rhs_index = StridedIndex::new(shape.dims(), rhs_stride);
let data = lhs_index
.zip(rhs_index)
.map(|(lhs_i, rhs_i)| lhs[lhs_i] * rhs[rhs_i])
.collect();
Ok(Storage::Cpu(CpuStorage::F32(data)))
}
(CpuStorage::F64(lhs), CpuStorage::F64(rhs)) => {
let lhs_index = StridedIndex::new(shape.dims(), lhs_stride);
let rhs_index = StridedIndex::new(shape.dims(), rhs_stride);
let data = lhs_index
.zip(rhs_index)
.map(|(lhs_i, rhs_i)| lhs[lhs_i] * rhs[rhs_i])
.collect();
Ok(Storage::Cpu(CpuStorage::F64(data)))
}
_ => {
// This should be covered by the dtype check above.
Err(Error::DTypeMismatchBinaryOp {
lhs: lhs.dtype(),
rhs: rhs.dtype(),
op: "add",
})
}
},
}
}
}