1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

use super::Tensor;
use crate::TchError;

pub struct Iter<T> {
    index: i64,
    len: i64,
    content: Tensor,
    phantom: std::marker::PhantomData<T>,
}

impl Tensor {
    pub fn iter<T>(&self) -> Result<Iter<T>, TchError> {
        Ok(Iter {
            index: 0,
            len: self.size1()?,
            content: self.shallow_clone(),
            phantom: std::marker::PhantomData,
        })
    }
}

impl std::iter::Iterator for Iter<i64> {
    type Item = i64;
    fn next(&mut self) -> Option<Self::Item> {
        if self.index >= self.len {
            return None;
        }
        let v = self.content.int64_value(&[self.index]);
        self.index += 1;
        Some(v)
    }
}

impl std::iter::Iterator for Iter<f64> {
    type Item = f64;
    fn next(&mut self) -> Option<Self::Item> {
        if self.index >= self.len {
            return None;
        }
        let v = self.content.double_value(&[self.index]);
        self.index += 1;
        Some(v)
    }
}

impl std::iter::Sum for Tensor {
    fn sum<I: Iterator<Item = Tensor>>(mut iter: I) -> Tensor {
        match iter.next() {
            None => Tensor::from(0.),
            Some(t) => iter.fold(t, |acc, x| x + acc),
        }
    }
}

impl<'a> std::iter::Sum<&'a Tensor> for Tensor {
    fn sum<I: Iterator<Item = &'a Tensor>>(mut iter: I) -> Tensor {
        match iter.next() {
            None => Tensor::from(0.),
            Some(t) => iter.fold(t.shallow_clone(), |acc, x| x + acc),
        }
    }
}