dfdx 0.13.0

Ergonomic auto differentiation in Rust, with pytorch like apis.
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
114
115
116
117
118
119

use super::*;
use std::ops::{Bound, RangeBounds};

fn get_start_bound(bound: Bound<&usize>) -> usize {
    match bound {
        Bound::Included(x) => *x,
        Bound::Excluded(x) => x + 1,
        Bound::Unbounded => 0,
    }
}

fn get_end_bound(bound: Bound<&usize>, size: usize) -> usize {
    match bound {
        Bound::Excluded(x) => *x,
        Bound::Included(x) => x + 1,
        Bound::Unbounded => size,
    }
}

pub trait SliceDim<R: RangeBounds<usize>>: Dim {
    type Sliced: Dim;

    fn slice(&self, range: &R) -> Option<Self::Sliced> {
        let size = self.size();

        let start_bound = get_start_bound(range.start_bound());
        let end_bound = get_end_bound(range.end_bound(), size);

        (end_bound <= size && start_bound <= end_bound)
            .then_some(end_bound - start_bound)
            .and_then(Self::Sliced::from_size)
    }
}

macro_rules! slice_dim_to_usize {
    ($range:ty) => {
        impl<D: Dim> SliceDim<$range> for D {
            type Sliced = usize;
        }
    };
}

slice_dim_to_usize!(std::ops::Range<usize>);
slice_dim_to_usize!(std::ops::RangeTo<usize>);
slice_dim_to_usize!(std::ops::RangeFrom<usize>);
slice_dim_to_usize!(std::ops::RangeInclusive<usize>);
slice_dim_to_usize!(std::ops::RangeToInclusive<usize>);

impl<D: Dim> SliceDim<std::ops::RangeFull> for D {
    type Sliced = D;

    fn slice(&self, _: &std::ops::RangeFull) -> Option<D> {
        Some(*self)
    }
}

pub trait SliceShape<R>: Shape {
    type Sliced: Shape<Concrete = Self::Concrete>;

    fn slice(&self, range: &R) -> Option<Self::Sliced>;
    fn first_idx_in_slice(&self, range: &R) -> usize;
}

impl SliceShape<()> for () {
    type Sliced = Self;

    fn slice(&self, _range: &()) -> Option<Self> {
        Some(())
    }

    fn first_idx_in_slice(&self, _range: &()) -> usize {
        0
    }
}

use super::broadcasts::length;

macro_rules! slice_shape {
    ([$($dim:ident)*] [$($range:ident)*] [$($idx:tt)*]) => {
        impl<$($dim: Dim),*, $($range: RangeBounds<usize>),*> SliceShape<($($range,)*)> for ($($dim,)*)
        where
            $($dim: SliceDim<$range>),*
        {
            type Sliced = ($($dim::Sliced,)*);

            fn slice(&self, range: &($($range,)*)) -> Option<Self::Sliced> {
                Some(($(self.$idx.slice(&range.$idx)?,)*))
            }

            fn first_idx_in_slice(&self, range: &($($range,)*)) -> usize {
                let strides = self.strides();
                $(get_start_bound(range.$idx.start_bound()) * strides[$idx] + )* 0
            }
        }

        impl<$($range: RangeBounds<usize>),*> SliceShape<($($range,)*)> for [usize; {length!($($range)*)}]
        where
            $(usize: SliceDim<$range>),*
        {
            type Sliced = ($(<usize as SliceDim<$range>>::Sliced,)*);

            fn slice(&self, range: &($($range,)*)) -> Option<Self::Sliced> {
                Some(($(self[$idx].slice(&range.$idx)?,)*))
            }

            fn first_idx_in_slice(&self, range: &($($range,)*)) -> usize {
                let strides = self.strides();
                $(get_start_bound(range.$idx.start_bound()) * strides[$idx] + )* 0
            }
        }
    }
}

slice_shape!([D1][R1][0]);
slice_shape!([D1 D2] [R1 R2] [0 1]);
slice_shape!([D1 D2 D3] [R1 R2 R3] [0 1 2]);
slice_shape!([D1 D2 D3 D4] [R1 R2 R3 R4] [0 1 2 3]);
slice_shape!([D1 D2 D3 D4 D5] [R1 R2 R3 R4 R5] [0 1 2 3 4]);
slice_shape!([D1 D2 D3 D4 D5 D6] [R1 R2 R3 R4 R5 R6] [0 1 2 3 4 5]);