topos 0.13.1

An autodiff compiler stack 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

use smallvec::smallvec;

use crate::{Element, Recordable, Shape, Tensor};

use super::{Cotangents, Operation, Reads, unary};

/// The `(count, size)` window pair at `axis`, `axis + 1` folded back
/// onto an axis of `extent`: [`Unfold`](super::Unfold)'s adjoint,
/// with the output-centric deterministic semantics of
/// [`Recordable::fold`].
///
/// It exists as an opcode because `unfold`'s derivative rule speaks
/// `fold`, so recorded gradients of windowed values need it on the
/// tape — the adjoint pairing the convolution design specified as
/// opcode-ready. The gradient of the operand is the incoming gradient
/// unfolded by the same parameters: the pair is self-adjoint in both
/// directions.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct Fold {
    pub(crate) axis: usize,
    pub(crate) size: usize,
    pub(crate) step: usize,
    pub(crate) dilation: usize,
    pub(crate) extent: usize,
}

impl Fold {
    /// Returns the arity: one operand.
    pub(crate) fn arity(&self) -> usize {
        1
    }

    /// Returns the read set of the derivative rule below.
    /// It reads no payloads: the cotangent unfolds by the parameters.
    pub(crate) fn reads(&self) -> Reads {
        Reads::NOTHING
    }

    /// Infers the result shape: the `(count, size)` pair at `axis`
    /// replaced by `extent`, requiring the operand pair to be exactly
    /// what unfolding an `extent` axis by these parameters produces.
    pub(crate) fn infer_shape(&self, operands: &[Shape]) -> Shape {
        let operand = unary(operands);
        assert!(
            self.axis + 1 < operand.rank(),
            "fold needs a (count, size) pair at axis {}, but {operand} has no pair there",
            self.axis
        );
        assert!(self.size > 0, "fold windows must hold at least one element");
        assert!(self.step > 0, "fold step must be positive");
        assert!(self.dilation > 0, "fold dilation must be positive");
        let span = self
            .dilation
            .checked_mul(self.size - 1)
            .and_then(|reach| reach.checked_add(1))
            .expect("fold window span overflows `usize`");
        assert!(
            span <= self.extent,
            "fold window span {span} exceeds the target extent {}",
            self.extent
        );
        let count = (self.extent - span) / self.step + 1;
        assert_eq!(
            operand.axes()[self.axis],
            count,
            "fold expects {count} windows at axis {} for extent {}, got {}",
            self.axis,
            self.extent,
            operand.axes()[self.axis]
        );
        assert_eq!(
            operand.axes()[self.axis + 1],
            self.size,
            "fold expects window size {} at axis {}, got {}",
            self.size,
            self.axis + 1,
            operand.axes()[self.axis + 1]
        );
        let mut folded: Vec<usize> = operand.axes().to_vec();
        folded[self.axis] = self.extent;
        folded.remove(self.axis + 1);
        Shape::new(folded)
    }
}

impl Fold {
    pub(crate) fn forward<E: Element>(&self, operands: &[&Tensor<E>]) -> Tensor<E> {
        unary(operands).fold(self.axis, self.size, self.step, self.dilation, self.extent)
    }
}

impl<Rule: Recordable> Operation<Rule> for Fold {
    fn backward(&self, _operands: &[&Rule], _output: &Rule, gradient: &Rule) -> Cotangents<Rule> {
        smallvec![Some(gradient.unfold(
            self.axis,
            self.size,
            self.step,
            self.dilation
        ))]
    }
}

#[cfg(test)]
#[path = "tests/fold_tests.rs"]
mod tests;