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

Fix the cpu implementation for narrow.

Browse Source
This commit is contained in:
laurent
2023-06-24 15:01:32 +01:00
parent 1b5f892d73
commit 0f34738831
2 changed files with 11 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
src/tensor.rs
View File

@ -316,18 +316,21 @@ impl Tensor {
} }
let mut dims = dims.to_vec(); let mut dims = dims.to_vec();
dims[dim] = length; dims[dim] = length;
let shape = Shape::from(dims); let adjusted_shape = Shape::from(dims);
let mut storage = self.device().zeros(&shape, self.dtype())?; let mut storage = self.device().zeros(&adjusted_shape, self.dtype())?;
let src_offset = self.stride[dim] * start; self.storage.copy_strided_src(
// TODO: This is incorrect, see the currently wrong test in tensor_tests.rs &mut storage,
self.storage /* dst_offset= */ 0,
.copy_strided_src(&mut storage, 0, &self.shape, &self.stride, src_offset)?; &adjusted_shape,
&self.stride,
/* src_offest= */ self.stride[dim] * start,
)?;
let op = if self.track_op() { let op = if self.track_op() {
Some(Op::Narrow(self.clone(), dim, start, length)) Some(Op::Narrow(self.clone(), dim, start, length))
} else { } else {
None None
}; };
Ok(from_storage(storage, shape, op, false)) Ok(from_storage(storage, adjusted_shape, op, false))
} }
pub fn softmax(&self, dim: usize) -> Result<Self> { pub fn softmax(&self, dim: usize) -> Result<Self> {
@ -881,7 +884,6 @@ impl Tensor {
pub fn backward(&self) -> Result<GradStore> { pub fn backward(&self) -> Result<GradStore> {
let sorted_nodes = self.sorted_nodes(); let sorted_nodes = self.sorted_nodes();
println!("{}", sorted_nodes.len());
let mut grads = GradStore::new(); let mut grads = GradStore::new();
grads.insert(self, self.ones_like()?); grads.insert(self, self.ones_like()?);
for node in sorted_nodes.iter() { for node in sorted_nodes.iter() {

3
tests/tensor_tests.rs
View File

@ -114,8 +114,7 @@ fn narrow() -> Result<()> {
let tensor = Tensor::new(data, &Device::Cpu)?; let tensor = Tensor::new(data, &Device::Cpu)?;
assert_eq!( assert_eq!(
tensor.narrow(2, 1, 2)?.to_vec3::<f32>()?, tensor.narrow(2, 1, 2)?.to_vec3::<f32>()?,
// TODO: this is broken at the moment! &[[[1.0, 4.0], [5.0, 9.0]], [[1.0, 7.0], [2.0, 8.0]]],
&[[[1., 4.], [1., 5.]], [[9., 2.], [1., 7.]]]
); );
Ok(()) Ok(())
} }