fastnn 0.3.3

A robust deep learning library in Rust with CUDA kernels
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

//! Derivatives of the rearranging ops, each the inverse rearrangement.

use crate::autograd::Backward;
use crate::tensor::Tensor;
use crate::tensor::shape;

use super::reduce_to;

/// Reshape the gradient back to the original shape.
pub struct ReshapeBackward {
    pub shape: Vec<usize>,
}

impl Backward for ReshapeBackward {
    fn backward(&self, grad: &Tensor) -> Vec<Tensor> {
        let dims: Vec<i64> = self.shape.iter().map(|&d| d as i64).collect();
        vec![grad.reshape(&dims)]
    }
    fn name(&self) -> &'static str {
        "Reshape"
    }
}

/// Apply the inverse permutation.
pub struct PermuteBackward {
    pub order: Vec<usize>,
}

impl Backward for PermuteBackward {
    fn backward(&self, grad: &Tensor) -> Vec<Tensor> {
        let mut inverse = vec![0usize; self.order.len()];
        for (position, &dim) in self.order.iter().enumerate() {
            inverse[dim] = position;
        }
        vec![grad.permute(&inverse)]
    }
    fn name(&self) -> &'static str {
        "Permute"
    }
}

/// The gradient of a copy is the sum over the copies.
pub struct ExpandBackward {
    pub shape: Vec<usize>,
}

impl Backward for ExpandBackward {
    fn backward(&self, grad: &Tensor) -> Vec<Tensor> {
        vec![reduce_to(grad, &self.shape)]
    }
    fn name(&self) -> &'static str {
        "Expand"
    }
}

/// Elements outside the slice get zero.
pub struct NarrowBackward {
    pub shape: Vec<usize>,
    pub axis: usize,
    pub start: usize,
}

impl Backward for NarrowBackward {
    fn backward(&self, grad: &Tensor) -> Vec<Tensor> {
        let (outer, size, inner) = shape::split_at_axis(&self.shape, self.axis);
        let taken = grad.dim(self.axis);
        let g = grad.to_vec();

        let mut out = vec![0.0f32; outer * size * inner];
        for o in 0..outer {
            let dst = (o * size + self.start) * inner;
            let src = o * taken * inner;
            out[dst..dst + taken * inner].copy_from_slice(&g[src..src + taken * inner]);
        }
        vec![Tensor::from_vec(out, &self.shape).to(grad.device())]
    }
    fn name(&self) -> &'static str {
        "Narrow"
    }
}

/// Split the gradient back into the pieces that were joined.
pub struct CatBackward {
    pub axis: usize,
    pub sizes: Vec<usize>,
}

impl Backward for CatBackward {
    fn backward(&self, grad: &Tensor) -> Vec<Tensor> {
        let mut start = 0;
        self.sizes
            .iter()
            .map(|&size| {
                let piece = grad.narrow(self.axis, start, size);
                start += size;
                piece
            })
            .collect()
    }
    fn name(&self) -> &'static str {
        "Cat"
    }
}