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

fix: fix index_select cuda kernel for src target dim different than ids dim when selecting dim > 0 (#1037)

Browse Source 
* fix: fix index_select cuda kernel for src target dim different than ids dim when selecting dim > 0

* cargo fmt
This commit is contained in:
Gonzalo
2023-10-05 14:46:13 -03:00
committed by GitHub
parent f0c619a4af
commit 8f7973958c
3 changed files with 21 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
candle-core/tests/tensor_tests.rs
View File

@ -680,6 +680,15 @@ fn index_select(device: &Device) -> Result<()> {
[3.0, 4.0, 5.0],
]
);
// Test when selecting dim > 0 with ids size different from elem count of
// target dim in source/input.
let ids = Tensor::new(&[1u32, 0u32, 1u32], device)?;
let t = Tensor::arange(1f32, 5f32, device)?.reshape((2, 2))?;
assert_eq!(t.to_vec2::<f32>()?, &[[1.0, 2.0], [3.0, 4.0]]);
let hs = t.index_select(&ids, 1)?;
assert_eq!(hs.to_vec2::<f32>()?, &[[2.0, 1.0, 2.0], [4.0, 3.0, 4.0]]);
Ok(())
}