aegir 2.0.0

Strongly-typed, reverse-mode autodiff library in 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

use super::*;
use std::ops::{Index, IndexMut};

/// Fixed rank, variable `DIM`-dimensional shape.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct SDynamic<const DIM: usize>(pub [usize; DIM]);

impl<const DIM: usize> Index<usize> for SDynamic<DIM> {
    type Output = usize;

    fn index(&self, idx: usize) -> &usize { self.0.index(idx) }
}

impl<const DIM: usize> IndexMut<usize> for SDynamic<DIM> {
    fn index_mut(&mut self, idx: usize) -> &mut usize { self.0.index_mut(idx) }
}

impl<const DIM: usize> Shape for SDynamic<DIM> {
    type Index = [usize; DIM];
    type IndexIter = multi_product::MultiProduct<DIM>;

    const DIM: usize = DIM;

    fn contains(&self, ix: [usize; DIM]) -> bool {
        ix.iter().zip(self.0.iter()).all(|(l, r)| l < r)
    }

    fn cardinality(&self) -> usize { self.0.iter().product() }

    fn indices(&self) -> Self::IndexIter { multi_product::MultiProduct::new(self.0) }
}

impl<const DIM: usize> Broadcast for SDynamic<DIM> {
    type Shape = SDynamic<DIM>;

    #[inline]
    fn broadcast(self, rhs: Self) -> Result<Self, IncompatibleShapes<Self>> {
        if self == rhs {
            Ok(self)
        } else {
            Err(IncompatibleShapes {
                left: self,
                right: rhs,
            })
        }
    }
}

impl<const DIM: usize> Broadcast<S0> for SDynamic<DIM> {
    type Shape = SDynamic<DIM>;

    #[inline]
    fn broadcast(self, _: S0) -> Result<Self, IncompatibleShapes<Self, S0>> { Ok(self) }
}

impl<const DIM: usize> Broadcast<SDynamic<DIM>> for S0 {
    type Shape = SDynamic<DIM>;

    #[inline]
    fn broadcast(self, rhs: SDynamic<DIM>) -> Result<Self::Shape, IncompatibleShapes<S0, SDynamic<DIM>>> {
        Ok(rhs)
    }
}

impl<const DIM: usize> std::fmt::Display for SDynamic<DIM> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "SDynamic({:?})", self.0)
    }
}

macro_rules! impl_add_dim {
    ($n:literal + $m:literal) => {
        impl Concat<SDynamic<$m>> for SDynamic<$n> {
            type Shape = SDynamic<{$n + $m}>;

            fn concat(self: SDynamic<$n>, rhs: SDynamic<$m>) -> Self::Shape {
                SDynamic(concat_arrays!(self.0, rhs.0))
            }

            fn concat_indices(left: [usize; $n], rhs: [usize; $m]) -> IndexOf<Self::Shape> {
                concat_arrays!(left, rhs)
            }
        }
    };
    ([$($n:literal),*] + $m:literal) => {
        $(impl_add_dim!($n + $m);)*
    };
    ($ns:tt + [$($m:literal),*]) => {
        $(impl_add_dim!($ns + $m);)*
    }
}

impl_add_dim!([0, 1, 2, 3, 4, 5, 6] + [0, 1, 2, 3, 4, 5, 6]);

// macro_rules! impl_drop_dim {
// ($n:literal) => {
// impl DropDim for SDynamic<$n> {
// type Lower = SDynamic<{$n - 1}>;

// fn drop_dim(self, dim: usize) -> SDynamic<{$n - 1}> {
// if dim >= Self::DIM { panic!("Lame..."); }

// SDynamic(array_init::array_init(|i| if i >= dim { self.0[i + 1] } else {
// self.0[i] })) }
// }
// };
// ($($n:literal),*) => {
// $(impl_drop_dim!($n);)*
// };
// }

// impl_drop_dim!(1, 2, 3, 4, 5, 6);