topos/facade/neural/
mlp.rs

1use static_assertions::assert_impl_all;
2
3use crate::{Element, Shape, Symbol, Tape, Tensor, Value};
4
5use super::{Activation, Linear, Module, Segment, Visitor};
6
7// Entry-time thread-safety contract; the anchor rationale is documented
8// in `network.rs`.
9assert_impl_all!(Mlp<f64>: Send, Sync);
10
11/// A multilayer perceptron: affine stages chained by topology, the
12/// convenience constructor over [`Linear`] and [`Activation`].
13///
14/// A topology such as `[3, 4, 4, 1]` defines three stages that map a
15/// `[batch, 3]` input to a `[batch, 1]` output. Hidden stages apply
16/// the caller's [`Activation`] after their affine transform; the
17/// output stage is affine alone. The contained stages retain
18/// parameter [`Symbol`]s, so the perceptron records in each
19/// compatible generation.
20#[derive(Debug, Clone)]
21pub struct Mlp<E> {
22    stages: Vec<Linear<E>>,
23    activation: Activation,
24}
25
26impl<E: Element> Mlp<E> {
27    /// Allocates the perceptron's stages on `tape` and returns it.
28    ///
29    /// `sizes` lists the value widths from the input width to the
30    /// output width. `activation` is applied after every hidden
31    /// stage's affine transform — caller-owned, like every
32    /// hyperparameter, with no default. `initializer` produces the
33    /// initial payload for each parameter from its shape — `[inputs,
34    /// outputs]` weights and `[outputs]` biases, stage by stage. The
35    /// initializer is responsible for returning payloads with the
36    /// requested shapes, and callers control details such as fan-in
37    /// scaling, randomness, and symmetry breaking.
38    ///
39    /// # Panics
40    /// Panics if `sizes` has fewer than two entries. It also propagates
41    /// [`Linear::new`] validation failures if initialized weights and
42    /// biases do not form valid parameter shapes.
43    pub fn new(
44        tape: &Tape<E>,
45        sizes: &[usize],
46        activation: Activation,
47        mut initializer: impl FnMut(&Shape) -> Tensor<E>,
48    ) -> Self {
49        assert!(
50            sizes.len() >= 2,
51            "an MLP topology needs an input and an output width"
52        );
53        let stages = sizes
54            .windows(2)
55            .map(|pair| {
56                let weights = initializer(&Shape::new([pair[0], pair[1]]));
57                let bias = initializer(&Shape::new([pair[1]]));
58                Linear::new(tape, weights, bias)
59            })
60            .collect();
61        Self { stages, activation }
62    }
63
64    /// Returns the symbols of all parameters, stage by stage: each
65    /// stage's weights, then its bias.
66    pub fn parameters(&self) -> impl Iterator<Item = Symbol> + '_ {
67        super::parameters(self).into_iter()
68    }
69}
70
71impl<E: Element> Module<E> for Mlp<E> {
72    /// Records the perceptron's expression over the `[batch, inputs]`
73    /// value `input` and returns the `[batch, outputs]` output value.
74    ///
75    /// # Panics
76    /// Panics if the parameters or `input` are not allocated on the
77    /// input's tape, or if `input` and the initialized stage shapes
78    /// are incompatible.
79    fn express<'tape>(&self, input: Value<'tape, E>) -> Value<'tape, E> {
80        let last = self.stages.len() - 1;
81        self.stages
82            .iter()
83            .enumerate()
84            .fold(input, |value, (index, stage)| {
85                let affine = stage.express(value);
86                if index == last {
87                    affine
88                } else {
89                    self.activation.express(affine)
90                }
91            })
92    }
93
94    fn visit(&self, visitor: &mut dyn Visitor) {
95        for (index, stage) in self.stages.iter().enumerate() {
96            visitor.enter(Segment::Index(index));
97            stage.visit(visitor);
98            visitor.leave();
99        }
100    }
101}
102
103#[cfg(test)]
104#[path = "tests/mlp_tests.rs"]
105mod tests;
topos/facade/neural/mlp.rs

topos/facade/neural/
mlp.rs
