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

use smallvec::smallvec;

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

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

/// The sliding windows of a value along one axis: the axis becomes a
/// `(count, size)` pair where window `w` starts at `w * step` and takes
/// every `dilation`-th element.
///
/// The forward is the [`Recordable::unfold`] strided view; the gradient
/// of the operand is the incoming gradient with every window
/// contribution summed back onto its source position, which is what
/// [`Recordable::fold`] computes — the two operations are adjoint. The
/// `dilation` slot ships ahead of any consumer so dilated convolution
/// later is a formula parameter, not a variant change.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct Unfold {
    pub(crate) axis: usize,
    pub(crate) size: usize,
    pub(crate) step: usize,
    pub(crate) dilation: usize,
}

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

    /// Returns the read set of the derivative rule below.
    /// It reads its operand for shape only, which a placeholder answers.
    pub(crate) fn reads(&self) -> Reads {
        Reads::NOTHING
    }

    /// Infers the result shape: `axis` replaced by the `(count, size)`
    /// window pair, requiring nonzero parameters and a dilated window
    /// span within the axis extent.
    pub(crate) fn infer_shape(&self, operands: &[Shape]) -> Shape {
        let operand = unary(operands);
        assert!(
            self.axis < operand.rank(),
            "unfold axis {} is out of rank for {operand}",
            self.axis
        );
        assert!(
            self.size > 0,
            "unfold windows must hold at least one element"
        );
        assert!(self.step > 0, "unfold step must be positive");
        assert!(self.dilation > 0, "unfold dilation must be positive");
        let extent = operand.axes()[self.axis];
        let span = self
            .dilation
            .checked_mul(self.size - 1)
            .and_then(|reach| reach.checked_add(1))
            .expect("unfold window span overflows `usize`");
        assert!(
            span <= extent,
            "unfold window span {span} exceeds axis {} extent {extent}",
            self.axis
        );
        let count = (extent - span) / self.step + 1;
        let mut unfolded: Vec<usize> = operand.axes().to_vec();
        unfolded[self.axis] = count;
        unfolded.insert(self.axis + 1, self.size);
        Shape::new(unfolded)
    }
}

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

impl<Rule: Recordable> Operation<Rule> for Unfold {
    fn backward(&self, operands: &[&Rule], _output: &Rule, gradient: &Rule) -> Cotangents<Rule> {
        let &operand = unary(operands);
        let extent = operand.shape().axes()[self.axis];
        smallvec![Some(gradient.fold(
            self.axis,
            self.size,
            self.step,
            self.dilation,
            extent
        ))]
    }
}