topos 0.12.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
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

use std::marker::PhantomData;

use static_assertions::assert_impl_all;

use crate::{Element, Symbol, Tape, Tensor, Value};

use super::{Module, Visitor};

// Entry-time thread-safety contract; the anchor rationale is documented
// in `network.rs`.
assert_impl_all!(Linear<f64>: Send, Sync);

/// The affine transform `input.matmul(weights) + bias`, unfused: an
/// activation is its own composition stage, which unlocks the
/// orderings a bundled activation forbids (pre-norm blocks,
/// activation-before-projection).
///
/// The weights are one `[inputs, outputs]` parameter and the bias is
/// one `[outputs]` parameter. The bias is broadcast explicitly across
/// the batch axis, so expressing the transform records a small, fixed
/// number of graph nodes regardless of parameter count. Parameters
/// are stored as [`Symbol`]s and resolved when the module records on
/// the family's [`Tape`].
#[derive(Debug, Clone)]
pub struct Linear<E> {
    weights: Symbol,
    bias: Symbol,
    _marker: PhantomData<E>,
}

impl<E: Element> Linear<E> {
    /// Allocates the transform's parameters on `tape` from their
    /// initial payloads and returns the module.
    ///
    /// The shapes are taken from the payloads: `weights` must be a
    /// rank-2 `[inputs, outputs]` payload and `bias` a rank-1
    /// `[outputs]` payload agreeing on `outputs`. Callers own
    /// initialization (fan-in scaling, randomness); the module records
    /// whatever it is given.
    ///
    /// # Panics
    /// Panics if `weights` is not rank 2, `bias` is not rank 1, or the
    /// two disagree on the number of outputs.
    pub fn new(tape: &Tape<E>, weights: Tensor<E>, bias: Tensor<E>) -> Self {
        let weights_shape = weights.shape();
        let bias_shape = bias.shape();
        assert_eq!(
            weights_shape.rank(),
            2,
            "linear weights must be rank 2, got {weights_shape}"
        );
        assert_eq!(
            bias_shape.rank(),
            1,
            "linear bias must be rank 1, got {bias_shape}"
        );
        assert_eq!(
            weights_shape.axes()[1],
            bias_shape.axes()[0],
            "linear weights {weights_shape} and bias {bias_shape} disagree on outputs"
        );
        Self {
            weights: tape.parameter(weights).symbol(),
            bias: tape.parameter(bias).symbol(),
            _marker: PhantomData,
        }
    }

    /// Returns the symbol of the `[inputs, outputs]` weight matrix.
    pub fn weights(&self) -> Symbol {
        self.weights
    }

    /// Returns the symbol of the `[outputs]` bias vector.
    pub fn bias(&self) -> Symbol {
        self.bias
    }

    /// Returns the symbols of the transform's parameters: the weights,
    /// then the bias.
    pub fn parameters(&self) -> impl Iterator<Item = Symbol> + '_ {
        super::parameters(self).into_iter()
    }
}

impl<E: Element> Module<E> for Linear<E> {
    /// Records the transform over the `[batch, inputs]` value `input`
    /// and returns the `[batch, outputs]` output value.
    ///
    /// # Panics
    /// Panics if the parameters or `input` are not allocated on
    /// `tape`, or if `input` and the weights are not compatible
    /// rank-2 matrices.
    fn express<'tape>(&self, input: Value<'tape, E>) -> Value<'tape, E> {
        let tape = input.tape();
        let weights = tape.resolve(self.weights);
        let bias = tape.resolve(self.bias);
        let product = input.matmul(weights);
        // The bias is repeated across the batch axis; its gradient sums
        // back along the same axis, one contribution per sample.
        product + bias.broadcast_along_like(0, product)
    }

    fn visit(&self, visitor: &mut dyn Visitor) {
        visitor.parameter("weights", self.weights);
        visitor.parameter("bias", self.bias);
    }
}

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