acme-tensor 0.3.1

A comprehensive tensor library for Rust with support for automatic-differentiation.
Documentation 
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
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109

/*
    Appellation: iterator <mod>
    Contrib: FL03 <jo3mccain@icloud.com>
*/
use super::LayoutIter;
use crate::TensorBase;
use core::marker::PhantomData;
use core::ptr;

/// An immutable iterator of the elements of a [tensor](crate::tensor::TensorBase)
/// Elements are visited in order, matching the layout of the tensor.
pub struct Iter<'a, T> {
    inner: LayoutIter,
    ptr: *const T,
    tensor: TensorBase<&'a T>,
}

impl<'a, T> Iter<'a, T> {
    pub(crate) fn new(tensor: TensorBase<&'a T>) -> Self {
        Self {
            inner: tensor.layout().iter(),
            ptr: unsafe { *tensor.as_ptr() },
            tensor,
        }
    }
}

impl<'a, T> Iterator for Iter<'a, T> {
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        let pos = self.inner.next()?;
        let item = self.tensor.get_by_index(pos.index())?;
        self.ptr = ptr::from_ref(item);
        unsafe { self.ptr.as_ref() }
    }
}

impl<'a, T> DoubleEndedIterator for Iter<'a, T> {
    fn next_back(&mut self) -> Option<Self::Item> {
        let pos = self.inner.next_back()?;
        let item = self.tensor.get_by_index(pos.index())?;
        self.ptr = ptr::from_ref(item);
        unsafe { self.ptr.as_ref() }
    }
}

impl<'a, T> ExactSizeIterator for Iter<'a, T> {
    fn len(&self) -> usize {
        self.tensor.size()
    }
}

impl<'a, T> From<TensorBase<&'a T>> for Iter<'a, T> {
    fn from(tensor: TensorBase<&'a T>) -> Self {
        Self::new(tensor)
    }
}

impl<'a, T> From<&'a TensorBase<T>> for Iter<'a, T> {
    fn from(tensor: &'a TensorBase<T>) -> Self {
        Self::new(tensor.view())
    }
}

pub struct IterMut<'a, T: 'a> {
    inner: LayoutIter,
    ptr: *mut T,
    tensor: TensorBase<&'a mut T>,
    _marker: PhantomData<&'a mut T>,
}

impl<'a, T> IterMut<'a, T> {
    pub(crate) fn new(tensor: &'a mut TensorBase<T>) -> Self {
        Self {
            inner: tensor.layout().iter(),
            ptr: tensor.as_mut_ptr(),
            tensor: tensor.view_mut(),
            _marker: PhantomData,
        }
    }
}

impl<'a, T> Iterator for IterMut<'a, T> {
    type Item = &'a mut T;

    fn next(&mut self) -> Option<Self::Item> {
        let pos = self.inner.next()?;
        let elem = self.tensor.get_mut_by_index(pos.index())?;
        self.ptr = ptr::from_mut(elem);
        unsafe { self.ptr.as_mut() }
    }
}

impl<'a, T> DoubleEndedIterator for IterMut<'a, T> {
    fn next_back(&mut self) -> Option<Self::Item> {
        let pos = self.inner.next_back()?;
        let elem = self.tensor.get_mut_by_index(pos.index())?;

        self.ptr = ptr::from_mut(elem);
        unsafe { self.ptr.as_mut() }
    }
}

impl<'a, T> ExactSizeIterator for IterMut<'a, T> {
    fn len(&self) -> usize {
        self.tensor.size()
    }
}