lumen_core/tensor/
iter.rs

1use crate::{StorageIndices, StorageRef, WithDType};
2
3use super::Tensor;
4
5pub struct TensorIter<'a, T> {
6    indexes: StorageIndices<'a>,
7    storage: StorageRef<'a, T>,
8}
9
10impl<'a, T: WithDType> Iterator for TensorIter<'a, T> {
11    type Item = T;
12
13    fn next(&mut self) -> Option<T> {
14        let index = self.indexes.next()?;
15        return self.storage.get(index)
16    }
17}
18
19impl<T: WithDType> Tensor<T> {
20    pub fn iter(&self) -> crate::Result<TensorIter<T>> {
21        Ok(TensorIter {
22            indexes: self.0.layout.storage_indices(),
23            storage: self.storage_ref(0)?,
24        })
25    }
26}
27
28pub trait ResettableIterator: Iterator {
29    fn reset(&mut self);
30}
31
32impl<'a, T: WithDType> ResettableIterator for TensorIter<'a, T> {
33    fn reset(&mut self) {
34        self.indexes.reset();
35    }
36}
37
38impl<'a, T: WithDType> ExactSizeIterator for TensorIter<'a, T> {
39    fn len(&self) -> usize {
40        self.indexes.len()
41    }
42}
lumen_core/tensor/iter.rs

lumen_core/tensor/
iter.rs
