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

//! Spatial pooling as a composed formula over the sliding-window view.
//!
//! Max pooling rides the same two single-axis `unfold`s as convolution
//! and needs no reduce opcode of its own: the maximum is a left-biased
//! fold of the existing binary `maximum` over the window lanes, so
//! ties route their gradient deterministically to the earliest lane.

use crate::{Element, Value};

/// Records the square windows of a pooling operation and returns them
/// as `[batch, channels, out_height, out_width, size * size]` lanes.
///
/// It records the pooling head: two unfolds, the axis permutation,
/// and the lane-merging reshape (a copy, since the window view
/// overlaps for `stride < size`).
fn window_lanes<'tape, E: Element>(
    input: Value<'tape, E>,
    size: usize,
    stride: usize,
) -> Value<'tape, E> {
    let shape = input.shape();
    assert_eq!(
        shape.rank(),
        4,
        "pooling input must be rank 4 [batch, channels, height, width], got {shape}"
    );
    assert!(size > 0, "pooling windows must hold at least one element");
    assert!(stride > 0, "pooling stride must be positive");
    let windows = input.unfold(2, size, stride, 1).unfold(4, size, stride, 1);
    let windows_shape = windows.shape();
    let axes = windows_shape.axes();
    windows
        .permute([0, 1, 2, 4, 3, 5])
        .reshape([axes[0], axes[1], axes[2], axes[4], size * size])
}

/// Records the `size x size` max pooling of the `[batch, channels,
/// height, width]` value `input` with `stride` and returns the pooled
/// `[batch, channels, out_height, out_width]` value.
///
/// The window maximum is a left-biased fold of [`Value::maximum`] over
/// the window lanes in row-major window order, so a tie routes its
/// gradient to the earliest tied position — deterministic, like every
/// tie rule in the crate.
///
/// # Panics
/// Panics if `input` is not rank 4, `size` or `stride` is zero, or a
/// window does not fit the spatial extents.
pub fn max_pool<'tape, E: Element>(
    input: Value<'tape, E>,
    size: usize,
    stride: usize,
) -> Value<'tape, E> {
    let lanes = window_lanes(input, size, stride);
    let mut largest = lanes.narrow(4, 0, 1);
    for lane in 1..size * size {
        largest = largest.maximum(lanes.narrow(4, lane, 1));
    }
    largest.squeeze(4)
}

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