redstone-ml 0.0.0

High-performance Machine Learning, Auto-Differentiation and Tensor Algebra crate for Rust
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
110
111
112
113

use crate::iterator::collapse_contiguous::collapse_to_uniform_stride;
use crate::ndarray::MAX_DIMS;

#[non_exhaustive]
pub struct FlatIndexGenerator
{
    ndims: usize,
    shape: [usize; MAX_DIMS],
    stride: [usize; MAX_DIMS],

    size: usize,
    iterator_index: usize,

    indices: [usize; MAX_DIMS], // current index along each dimension
    flat_index: usize,
}

impl FlatIndexGenerator {
    pub(crate) fn from(shape: &[usize], stride: &[usize]) -> Self {
        let (shape, stride) = collapse_to_uniform_stride(shape, stride);
        let ndims = shape.len();
        let size = shape.iter().product();

        let mut new_shape = [0; MAX_DIMS];
        let mut new_stride = [0; MAX_DIMS];

        new_shape[0..ndims].copy_from_slice(&shape);
        new_stride[0..ndims].copy_from_slice(&stride);

        Self {
            ndims,
            shape: new_shape,
            stride: new_stride,
            size,
            iterator_index: 0,
            indices: [0; MAX_DIMS],
            flat_index: 0,
        }
    }

    #[inline(always)]
    pub fn advance_by(&mut self, mut n: usize) {
        let remaining = self.size - self.iterator_index;
        n = n.min(remaining);

        if n == 0 {
            return;
        }
        self.iterator_index += n;

        let mut carry = n;
        for i in (0..self.ndims).rev() {
            let dim = self.shape[i];
            let idx = &mut self.indices[i];

            let total = *idx + carry;
            *idx = total % dim;
            carry = total / dim;

            self.flat_index += self.stride[i] * (*idx - self.indices[i]);
        }
    }
}

impl Iterator for FlatIndexGenerator {
    type Item = usize;

    #[inline(always)]
    fn next(&mut self) -> Option<Self::Item> {
        if self.iterator_index == self.size {
            return None;
        }

        let return_index = self.flat_index;

        let mut i = self.ndims;
        while i > 0 {
            i -= 1;

            unsafe {
                let idx = self.indices.get_unchecked_mut(i);
                *idx += 1;

                if *idx < *self.shape.get_unchecked(i) {
                    self.flat_index += *self.stride.get_unchecked(i);
                    break;
                }

                *idx = 0; // reset this dimension and carry over to the next
                self.flat_index -= *self.stride.get_unchecked(i) * (*self.shape.get_unchecked(i) - 1);
            }
        }

        self.iterator_index += 1;
        Some(return_index)
    }
}

impl Clone for FlatIndexGenerator {
    fn clone(&self) -> Self {
        Self {
            ndims: self.ndims,
            shape: self.shape,
            stride: self.stride,

            size: self.size,
            iterator_index: self.iterator_index,

            indices: self.indices,
            flat_index: self.flat_index,
        }
    }
}