topos 0.13.0

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

use smallvec::{SmallVec, smallvec};

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

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

/// A permutation of a value's axes: axis `i` of the result takes axis
/// `order[i]` of the operand.
///
/// The gradient of the operand is the incoming gradient reordered by the
/// inverse permutation.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct Permute {
    pub(crate) order: SmallVec<[usize; 4]>,
}

impl Permute {
    /// 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 permutes back.
    pub(crate) fn reads(&self) -> Reads {
        Reads::NOTHING
    }

    /// Infers the result shape: the operand's axes reordered by `order`,
    /// which must be a permutation of the operand's axes.
    pub(crate) fn infer_shape(&self, operands: &[Shape]) -> Shape {
        let operand = unary(operands);
        assert_eq!(
            self.order.len(),
            operand.rank(),
            "permute order must cover every axis of {operand}"
        );
        let mut seen = vec![false; operand.rank()];
        for &axis in &self.order {
            assert!(
                axis < operand.rank(),
                "permute axis {axis} is out of rank for {operand}"
            );
            assert!(
                !std::mem::replace(&mut seen[axis], true),
                "permute order repeats axis {axis}"
            );
        }
        Shape::new(self.order.iter().map(|&axis| operand.axes()[axis]))
    }

    /// Returns the inverse permutation: the order that undoes `self.order`.
    fn inverse(&self) -> SmallVec<[usize; 4]> {
        let mut inverse: SmallVec<[usize; 4]> =
            std::iter::repeat_n(0usize, self.order.len()).collect();
        for (position, &axis) in self.order.iter().enumerate() {
            inverse[axis] = position;
        }
        inverse
    }
}

impl Permute {
    pub(crate) fn forward<E: Element>(&self, operands: &[&Tensor<E>]) -> Tensor<E> {
        unary(operands).permute(&self.order)
    }
}

impl<Rule: Recordable> Operation<Rule> for Permute {
    fn backward(&self, _operands: &[&Rule], _output: &Rule, gradient: &Rule) -> Cotangents<Rule> {
        smallvec![Some(gradient.permute(&self.inverse()))]
    }
}