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

// Trait which recursively flattens a vector
// [0, 50, 100] -> [0, 50, 100]
// [[0], [50, 50], 100] -> [0, 50, 50, 100]
// [[[0]]] -> [0]

use crate::recursive_trait_base_cases;
use crate::util::shape::Shape;
use std::fmt::Debug;

use std::ptr::copy_nonoverlapping;

pub(crate) trait Flatten<A> {
    fn flatten(self) -> Vec<A>;
}

impl<A, T> Flatten<A> for Vec<T>
where
    T: Flatten<A>,
{
    fn flatten(self) -> Vec<A> {
        self.into_iter()
            .flat_map(|nested| nested.flatten().into_iter())
            .collect()
    }
}

impl<A, T, const N: usize> Flatten<A> for [T; N]
where
    T: Flatten<A>,
    [T; N]: Shape,
    A: Debug,
{
    fn flatten(mut self) -> Vec<A> {
        assert!(align_of::<T>() >= align_of::<A>(), "alignment mismatch");  // ChatGPT suggested

        let len = self.shape().iter().product();
        let mut result = Vec::with_capacity(len);

        let src = self.as_mut_ptr() as *mut A;
        let dst = result.as_mut_ptr();

        unsafe {
            copy_nonoverlapping(src, dst, len);
            result.set_len(len);
        }
        result
    }
}

macro_rules! flatten_trait {
    ( $dtype:ty ) => {
        impl Flatten<$dtype> for Vec<$dtype> {
            fn flatten(self) -> Vec<$dtype> {
                self
            }
        }

        impl<const N: usize> Flatten<$dtype> for [$dtype; N] {
            fn flatten(self) -> Vec<$dtype> {
                Vec::from(self)
            }
        }
    };
}

recursive_trait_base_cases!(flatten_trait);