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

Move the StridedIndex in its own module.

Browse Source
This commit is contained in:
laurent
2023-06-21 07:44:36 +01:00
parent 23db8a7da8
commit 3a5405ca6d
4 changed files with 66 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

62
src/storage.rs
View File

@ -1,4 +1,4 @@
use crate::{DType, Device, Error, Result, Shape};
use crate::{DType, Device, Error, Result, Shape, StridedIndex};
// TODO: Think about whether we would be better off with a dtype and
// a buffer as an owned slice of bytes.
@ -17,66 +17,6 @@ impl CpuStorage {
}
}
#[derive(Debug)]
pub(crate) struct StridedIndex<'a> {
next_storage_index: Option<usize>,
multi_index: Vec<usize>,
dims: &'a [usize],
stride: &'a [usize],
}
impl<'a> StridedIndex<'a> {
pub(crate) fn new(dims: &'a [usize], stride: &'a [usize]) -> Self {
let elem_count: usize = dims.iter().product();
let next_storage_index = if elem_count == 0 {
None
} else {
// This applies to the scalar case.
Some(0)
};
StridedIndex {
next_storage_index,
multi_index: vec![0; dims.len()],
dims,
stride,
}
}
}
impl<'a> Iterator for StridedIndex<'a> {
type Item = usize;
fn next(&mut self) -> Option<Self::Item> {
let storage_index = match self.next_storage_index {
None => return None,
Some(storage_index) => storage_index,
};
let mut updated = false;
for (multi_i, max_i) in self.multi_index.iter_mut().zip(self.dims.iter()).rev() {
let next_i = *multi_i + 1;
if next_i < *max_i {
*multi_i = next_i;
updated = true;
break;
} else {
*multi_i = 0
}
}
self.next_storage_index = if updated {
let next_storage_index = self
.multi_index
.iter()
.zip(self.stride.iter())
.map(|(&x, &y)| x * y)
.sum();
Some(next_storage_index)
} else {
None
};
Some(storage_index)
}
}
#[derive(Debug, Clone)]
pub enum Storage {
Cpu(CpuStorage),