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

Add a currently wrong test for narrow.

Browse Source
This commit is contained in:
laurent
2023-06-24 08:50:37 +01:00
parent d6cb4f1c53
commit 1b5f892d73
2 changed files with 16 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
src/tensor.rs
View File

@ -300,7 +300,7 @@ impl Tensor {
/// Returns a new tensor that is a narrowed version of the input, the dimension `dim`
/// ranges from `start` to `start + length`.
// TODO: Once we've refactor the shape and strides, make this return a view of the same data
// TODO: Once we've refactored the shape and strides, make this return a view of the same data
// rather than copying.
pub fn narrow(&self, dim: usize, start: usize, length: usize) -> Result<Self> {
let dims = self.shape().dims();
@ -318,7 +318,8 @@ impl Tensor {
dims[dim] = length;
let shape = Shape::from(dims);
let mut storage = self.device().zeros(&shape, self.dtype())?;
let src_offset = 0; // TODO
let src_offset = self.stride[dim] * start;
// TODO: This is incorrect, see the currently wrong test in tensor_tests.rs
self.storage
.copy_strided_src(&mut storage, 0, &self.shape, &self.stride, src_offset)?;
let op = if self.track_op() {

13
tests/tensor_tests.rs
View File

@ -1,3 +1,4 @@
// TODO: Also test the cuda backend.
use candle::{DType, Device, Result, Tensor};
#[test]
@ -106,3 +107,15 @@ fn softmax() -> Result<()> {
);
Ok(())
}
#[test]
fn narrow() -> Result<()> {
let data = &[[[3f32, 1., 4.], [1., 5., 9.]], [[2., 1., 7.], [8., 2., 8.]]];
let tensor = Tensor::new(data, &Device::Cpu)?;
assert_eq!(
tensor.narrow(2, 1, 2)?.to_vec3::<f32>()?,
// TODO: this is broken at the moment!
&[[[1., 4.], [1., 5.]], [[9., 2.], [1., 7.]]]
);
Ok(())
}