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

use smallvec::{SmallVec, smallvec};

use crate::backend::Formula;

use super::batch_norm::BatchNormalization;
use super::reduce_window::ReduceWindow;
use super::window::WindowProduct;

/// A recognized pattern rooted at one plan node.
///
/// It is a compile-time match over frozen structure, not a tape
/// rewrite — and it carries no policy. What to do with a match belongs
/// to the consumers: each owns a repertoire (the patterns it can act
/// on) and an action per variant in its own module — the forward
/// run's kernel table lives beside the plan, the raises beside the
/// emitter. The shape mirrors `Op`, the role does not: a
/// `Op` carries its rules because they are the single spec,
/// while a pattern has as many interpretations as consumers.
#[derive(Debug, Clone)]
pub(crate) enum Pattern {
    /// Canonical im2col chain feeding a rank-2 `matmul`.
    WindowProduct(WindowProduct),
    /// Canonical max-pool window fold ending in the facade squeeze.
    ReduceWindow(ReduceWindow),
    /// Batch normalization by the batch's own statistics, with the
    /// mean and variance as named results.
    BatchNormTraining(BatchNormalization),
    /// Batch normalization by supplied statistics.
    BatchNormInference(BatchNormalization),
}

impl Pattern {
    /// The vocabulary entry this pattern is the graph face of; the
    /// consumers look its coverage up under this name.
    pub(crate) fn formula(&self) -> Formula {
        match self {
            Pattern::WindowProduct(_) => Formula::WindowProduct,
            Pattern::ReduceWindow(_) => Formula::ReduceWindow,
            Pattern::BatchNormTraining(_) => Formula::BatchNormTraining,
            Pattern::BatchNormInference(_) => Formula::BatchNormInference,
        }
    }

    /// The named-result slots a fusing consumer's action writes back
    /// at the root — materialized values, unlike unnamed interiors,
    /// which are never computed. Empty for single-result patterns.
    pub(crate) fn named(&self) -> SmallVec<[usize; 4]> {
        match self {
            Pattern::BatchNormTraining(group) => smallvec![group.mean, group.variance],
            Pattern::WindowProduct(_)
            | Pattern::ReduceWindow(_)
            | Pattern::BatchNormInference(_) => SmallVec::new(),
        }
    }
}