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

Support negative steps in arange. (#1218)

Browse Source
This commit is contained in:
Laurent Mazare
2023-10-30 08:40:54 +01:00
committed by GitHub
parent 174b208052
commit 5fc66bd4ba
2 changed files with 33 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
candle-core/src/tensor.rs
View File

@ -385,11 +385,21 @@ impl Tensor {
step: D,
device: &Device,
) -> Result<Self> {
if D::is_zero(&step) {
crate::bail!("step cannot be zero")
}
let mut data = vec![];
let mut current = start;
while current < end {
data.push(current);
current += step;
if step >= D::zero() {
while current < end {
data.push(current);
current += step;
}
} else {
while current > end {
data.push(current);
current += step;
}
}
let len = data.len();
Self::from_vec_impl(data, len, device, false)