3.6. Pooling Layers

Pooling downsamples a feature map. It summarizes each local window with a single number. The spatial dimensions shrink, and the channel count stays the same. Pooling is the cheap counterpart to the convolutions in 3.5. Convolutional Layers. It has no weights and no matrix multiply. It only reduces a sliding window to 1 number.

RustyML ships pooling as a family of 12 layers. All 12 layers use 1 engine that works for any spatial rank. This page covers the constructors, the output shapes, and the difference between max and average pooling. It also covers how gradients route backward, and what “no learnable parameters” means for summary() and model persistence. The last section covers the 3 upsampling layers. They run the same idea in reverse, and give a decoder its way back to the input resolution.

3.6.1. The family: 2 reductions, 3 ranks, plus global variants

Every pooling layer picks 1 of 2 reductions, max or average. It applies that reduction at 1 of 3 spatial ranks: 1D, 2D, or 3D. Each combination also has a “global” variant that collapses the whole spatial extent at once. That is 2 x 3 x 2 = 12 concrete types. All 12 types reduce to 4 functions in a private pooling_engine.

The engine derives the spatial rank at run time from input.ndim() - 2. 1 loop serves 1D, 2D, and 3D. The only per-layer differences are the reduction kind, and how the public pool_size/strides tuples flatten into slices.

Tensors follow the channels-last convention [batch, spatial..., channels], the same convention the convolution layers use. A 1D layer takes a 3D tensor, a 2D layer takes a 4D tensor, and a 3D layer takes a 5D tensor. Windowed pooling keeps the rank and shrinks the spatial axes. Global pooling drops the spatial axes and returns [batch, channels].

The channel axis sits innermost, so a whole position reduces at once. The engine works out the window geometry, the bounds checks, and the count of real elements a Same-padded average divides by, once per output position. Every channel shares that same geometry.

1 naming point differs by rank. The 1D windowed layers expose with_stride, a single usize. The 2D and 3D layers expose with_strides, a tuple. This matches the shape of pool_size at each rank.

3.6.2. Constructors, pool size, stride, and padding

Windowed layers take the pool window alone. For example: MaxPooling2D::new((2, 2)). The 6 windowed constructors are infallible now. Each returns the layer directly with no Result, because a pool window on its own cannot be wrong.

build checks the window against the input where the input arrives. 2 builder methods override the defaults. Both are optional:

• with_stride / with_strides sets the step between windows. If you never call it, the stride defaults to the pool size. This gives non-overlapping windows, which is the usual default. Pass a smaller stride for overlapping windows.
• with_padding sets PaddingType::Valid (the default, no padding) or PaddingType::Same. This is the same PaddingType enum the convolution layers use.

Global layers take no arguments at all: GlobalMaxPooling2D::new(). There is no window, stride, or padding to configure, because the window is the whole spatial plane.

// Windowed 2D max pooling, 3x3 window, stride 2, Same padding:
let layer = MaxPooling2D::new((3, 3))
    .with_strides((2, 2))
    .unwrap()
    .with_padding(PaddingType::Same);

// Global 2D average pooling needs nothing:
let head = GlobalAveragePooling2D::new();

build validates the shape that reaches the layer, and the model build supplies that shape from the layer in front. A window that does not fit is therefore refused before any data moves, and the message names the position of the layer and its type. A built layer then refuses an input that disagrees with the shape it was built for, on every axis except the batch axis. A plane smaller than the window can therefore no longer reach the geometry at all.

The table below lists the failure modes. See 1.6. Error Handling for the error taxonomy:

Every one of these arrives from the model build, wrapped in a message that names the position of the layer and its type. A window that does not fit its feature map used to surface as an unsigned underflow in the middle of a forward pass. It is now a Result that the build hands back, on the line that assembles the model.

3.6.3. Output-shape formulas

For Valid padding, each spatial axis shrinks by the same formula the convolution layers use. The division floors, because the engine drops the trailing remainder:

out = (in - pool) / stride + 1

For Same padding, the output rounds up to ceil(in / stride). The engine pads symmetrically, with the extra cell on the trailing edge. This matches the convolution engine. Global pooling ignores padding and window size. It always produces [batch, channels].

compute_output_shape returns these sizes, and it is pure. Every pooling layer answers it on the shape you hand it, built or not. summary() therefore prints a real output shape for a windowed layer and a global layer alike.

Layer::output_shape() is the other reader. It runs the same algebra against the shape the layer itself holds. It frees the batch axis first, because 1 built layer serves every batch size. A layer that no build has touched holds no shape, and it reports Unknown. Inside a model that difference is invisible, because the build gave every position its shape.

use ndarray::Array;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::{Layer, UnaryLayer};
use rustyml::neural_network::{Ctx, Shape};

fn main() {
    // `compute_output_shape` is pure, so an unbuilt layer answers for any input shape.
    let mp = MaxPooling2D::new((2, 2));
    let plane = Shape::known(&[1, 6, 6, 3]);
    println!("MaxPooling2D:     {}", mp.compute_output_shape(&plane).unwrap()); // (1, 3, 3, 3)

    let ap = AveragePooling1D::new(2);
    let line = Shape::known(&[1, 6, 1]);
    println!("AveragePooling1D: {}", ap.compute_output_shape(&line).unwrap()); // (1, 3, 1)

    // A global layer answers the same way, and it keeps a free batch axis free.
    let gap = GlobalAveragePooling2D::new();
    let images = Shape::with_free_batch(&[3, 5, 5, 4]);
    println!("GlobalAvgPool2D:  {}", gap.compute_output_shape(&images).unwrap()); // (None, 4)

    // `output_shape` reads the shape the layer itself holds, which a global layer only
    // learns from a build. `forward_mut` builds it from the tensor.
    let mut gap = GlobalAveragePooling2D::new();
    println!("before the pass:  {}", gap.output_shape()); // Unknown

    let x = Array::from_elem(ndarray::IxDyn(&[3, 5, 5, 4]), 1.0f32);
    let mut ctx = Ctx::inference();
    let out = gap.forward_mut(&x, &mut ctx).unwrap();
    assert_eq!(out.shape(), &[3, 4]);
    println!("after the pass:   {}", gap.output_shape()); // (None, 4)
}

Same padding is not only a shape adjustment. The padded cells are virtual. The forward and backward passes skip out-of-bounds positions instead of substituting zeros. This matters for averaging. An edge window divides by the count of real, in-bounds elements, not by the window area.

A pad cell therefore never enters the divisor. Same-padded average pooling does not dilute edges toward zero.

use ndarray::Array;
use rustyml::neural_network::Ctx;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::UnaryLayer;

fn main() {
    // 3x3 input, 2x2 window, stride 2. Valid padding drops the last row and column.
    // Same padding rounds the output up to ceil(3/2) = 2. The trailing windows see only
    // their in-bounds cells, because the padding is virtual.
    let x = Array::from_shape_vec((1, 3, 3, 1), (1..=9).map(|v| v as f32).collect())
        .unwrap()
        .into_dyn();
    let mut ctx = Ctx::inference();

    let mut max_same = MaxPooling2D::new((2, 2))
        .with_strides((2, 2))
        .unwrap()
        .with_padding(PaddingType::Same);
    let m = max_same.forward_mut(&x, &mut ctx).unwrap();
    assert_eq!(m.shape(), &[1, 2, 2, 1]);
    // [[max(1,2,4,5), max(3,6)], [max(7,8), 9]] = [[5, 6], [8, 9]]
    assert_eq!(m.iter().copied().collect::<Vec<_>>(), vec![5.0, 6.0, 8.0, 9.0]);

    let mut avg_same = AveragePooling2D::new((2, 2))
        .with_strides((2, 2))
        .unwrap()
        .with_padding(PaddingType::Same);
    let a = avg_same.forward_mut(&x, &mut ctx).unwrap();
    // Averages divide by the count of REAL cells, and not by the window area:
    // [[(1+2+4+5)/4, (3+6)/2], [(7+8)/2, 9/1]] = [[3.0, 4.5], [7.5, 9.0]]
    assert_eq!(a.iter().copied().collect::<Vec<_>>(), vec![3.0, 4.5, 7.5, 9.0]);
}

3.6.4. Max versus average: semantics and when each helps

Both reductions run over the same window. They differ only in what they keep. Max pooling records the single largest activation and discards the rest. It acts as a detector for “did this feature fire anywhere in the window”, and it tolerates translation. Average pooling keeps the mean. It preserves the overall magnitude of the region, and it smooths instead of selecting.

use ndarray::Array;
use rustyml::neural_network::Ctx;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::UnaryLayer;

fn main() {
    // The same 4x4 plane, pooled with a non-overlapping 2x2 window (stride defaults to 2).
    let data: Vec<f32> = (0..16).map(|v| v as f32).collect();
    let x = Array::from_shape_vec((1, 4, 4, 1), data).unwrap().into_dyn();
    let mut ctx = Ctx::inference();

    let mut max_pool = MaxPooling2D::new((2, 2));
    let m = max_pool.forward_mut(&x, &mut ctx).unwrap();
    // window maxima: [[5, 7], [13, 15]]
    assert_eq!(m.iter().copied().collect::<Vec<_>>(), vec![5.0, 7.0, 13.0, 15.0]);

    let mut avg_pool = AveragePooling2D::new((2, 2));
    let a = avg_pool.forward_mut(&x, &mut ctx).unwrap();
    // window means: [[2.5, 4.5], [10.5, 12.5]]
    assert_eq!(a.iter().copied().collect::<Vec<_>>(), vec![2.5, 4.5, 10.5, 12.5]);
}

Use max pooling inside the convolutional stack of a discriminative model. It keeps the strongest response through downsampling, and the exact position in the window does not matter. Use average pooling when magnitude matters more than peaks. Use the global average reduction at the end of a network. Averaging the whole plane gives a stable, smooth summary per channel.

2 engine details matter in edge cases. Max pooling breaks ties toward the first (lowest-index) maximum, using a strict > comparison. Max pooling also propagates NaN on purpose. Once a NaN enters a window, it wins and stays there. The layer therefore drops no NaN silently.

3.6.5. Gradient routing: backprop through pooling

Pooling has no parameters to update. It still must route the upstream gradient back to the inputs that produced its output. The 2 reductions route this gradient differently.

Max pooling is winner-takes-gradient. During the forward pass, each output records the flat index of the input element it picked. This index is the “arg-max”, and the pass parks it in the context. During backward, each upstream gradient scatters to that one position. Every other input in the window gets zero. When windows overlap and the same input wins more than 1 window, the contributions accumulate.

Average pooling instead spreads each output gradient evenly across its window. Every in-bounds element of the window receives grad / count. Overlaps accumulate here too.

Global max pooling routes each channel’s gradient to its single arg-max element. Global average pooling spreads each channel’s gradient evenly across the whole plane, as grad / spatial_size.

use ndarray::Array;
use rustyml::neural_network::Ctx;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::UnaryLayer;

fn main() {
    let data: Vec<f32> = (0..16).map(|v| v as f32).collect();
    let x = Array::from_shape_vec((1, 4, 4, 1), data).unwrap().into_dyn();
    let grad = Array::ones((1, 2, 2, 1)).into_dyn();

    // A training context holds what the backward pass takes back.
    let mut ctx = Ctx::training();

    // Max pooling: each window's gradient reaches only the winning cell (here the 4 maxima
    // at flat indices 5, 7, 13, 15). Every other cell receives 0.
    let mut max_pool = MaxPooling2D::new((2, 2));
    max_pool.forward_mut(&x, &mut ctx).unwrap();
    let gmax = max_pool.backward(&grad, &mut ctx).unwrap();
    let expected_max = vec![
        0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 1.0,
    ];
    assert_eq!(gmax.iter().copied().collect::<Vec<_>>(), expected_max);

    // Average pooling: each window's gradient is spread evenly over its cells (1.0 / 4 = 0.25).
    let mut avg_pool = AveragePooling2D::new((2, 2));
    avg_pool.forward_mut(&x, &mut ctx).unwrap();
    let gavg = avg_pool.backward(&grad, &mut ctx).unwrap();
    assert!(gavg.iter().all(|&g| (g - 0.25).abs() < 1e-6));
}

Because max pooling depends on the arg-max cache, the mode of the pass matters. A forward pass under Ctx::training() parks that cache in the context. A forward pass under Ctx::inference() parks nothing at all, which is what lets several threads run inference through 1 shared model. So backward works only against a context that a training forward pass wrote. A backward with a fresh context, or with a context that an inference pass left empty, returns Error::NeuralNetwork(NnError::ForwardPassNotRun(_)).

Inside Sequential::fit, this never happens, because training always runs the forward pass first. You trip this error only when you drive a bare layer by hand. Average pooling parks only the input shape, because it needs the geometry, and not the values. Global average pooling parks only the shape too. The “run forward before backward” contract is the same across all 12 layers.

3.6.6. Global pooling as a modern head

A traditional convolutional classifier ends with Flatten, followed by a large Dense layer. This design ties the network to 1 exact input resolution. It also puts most of the parameters into that last matrix.

Global average pooling replaces that head with a reduction that has no parameters. Network-in-Network and ResNet made this design common. The layer collapses each channel to its mean, so each channel yields 1 feature. A small Dense classifier follows the pooling layer.

The pooling step has no parameters, so it cannot overfit. This gives the head free regularization. The classifier’s input width depends only on the channel count, not on H x W.

This size independence is real for the pooling layer taken on its own. A GlobalAveragePooling2D that no build has touched accepts any height and width, but inside a model it does not stay that way. The build gives it the shape the convolution in front produces. A built layer refuses an input that disagrees with that shape on any axis except the batch axis. The upstream convolution refuses the same way, and the build sizes the Dense head from the pooled channel count. A built model therefore expects 1 consistent resolution, end to end, and it says so before any data moves.

The payoff of the global-pooling head is the parameter-free reduction, resistant to overfitting, and the tidy [N, C] -> Dense interface. It does not give automatic variable-resolution inference.

use rustyml::neural_network::Shape;
use rustyml::neural_network::sequential::SequentialBuilder;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::optimizers::*;
use rustyml::neural_network::losses::*;
use ndarray::Array;

fn main() {
    // A small classifier head: Conv -> pool -> global pool -> Dense.
    let x = Array::from_shape_fn((2, 8, 8, 1), |(b, i, j, _)| {
        (i + j) as f32 * 0.1 + b as f32
    })
    .into_dyn();
    let y = Array::ones((2, 3)).into_dyn();

    let mut model = SequentialBuilder::new()
        .add(Conv2D::new(4, (3, 3), (1, 1), Activation::ReLU).unwrap()) // -> [2, 6, 6, 4]
        .add(MaxPooling2D::new((2, 2)))                                 // -> [2, 3, 3, 4]
        .add(GlobalAveragePooling2D::new())                             // -> [2, 4]
        .add(Dense::new(3, Activation::ReLU).unwrap())                  // -> [2, 3]
        .build(&Shape::known(x.shape()))
        .unwrap();
    model.compile(RMSprop::new(0.001, 0.9, 1e-8, 0.0).unwrap(), MeanSquaredError::new());

    model.summary();
    model.fit(&x, &y, 2).unwrap();

    let prediction = model.predict(&x).unwrap();
    assert_eq!(prediction.shape(), &[2, 3]);
}

The summary() call above runs before fit, and every row already carries a real output shape. The global-pooling row is no exception. The build gave that layer the shape the pooling layer in front produces, and compute_output_shape collapses the 2 spatial axes of it. No tensor has flowed through the model yet:

Model: "sequential"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Layer (type)                    ┃ Output Shape           ┃       Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ conv2d (Conv2D)                 │ (2, 6, 6, 4)           │            40 │
│ maxpooling2d (MaxPooling2D)     │ (2, 3, 3, 4)           │             0 │
│ globalaveragepooling2d (GlobalAveragePooling2D) │ (2, 4)                 │             0 │
│ dense (Dense)                   │ (2, 3)                 │            15 │
└─────────────────────────────────┴────────────────────────┴───────────────┘
 Total params: 55 (220 B)
 Trainable params: 55 (220 B)
 Non-trainable params: 0 (0 B)

Both pooling rows show 0 in the Param # column. Neither one adds to any of the 3 totals. The shapes do not move after fit either, because they came from the build and not from a tensor.

3.6.7. No learnable parameters: what that means for persistence

Pooling layers report ParamCounts::none(), which is 0 trainable elements and 0 non-trainable ones. Their weights() list is empty, so they contribute no checkpoint path at all. LayerBase::parameters_mut gives back the empty vector too, so the optimizer finds nothing to update. The gradient store of the context holds nothing for them either. This is why they never appear in the trainable or non-trainable totals above. It is also why they add nothing to a saved file’s size.

Persistence is where “no parameters” has a real consequence. save_to_path writes 1 record per layer. That record holds the layer type name, the shape the layer was built for, and every named array of the layer. For a pooling layer, the array list is empty. So the file stores the type name and the build shape for that layer, and no array at all.

load_from_path does not reconstruct the architecture. You must rebuild the model with the same layers first, then load the weights. The load walks the layers position by position. It checks that the layer count matches. It also checks that each saved type string equals the type at the same index. It compares the build shape too, wherever the file and the layer both carry one.

Only then does it apply any weight. A pooling layer has no weight to restore, but it must still occupy its exact position, with its exact type. Otherwise the load fails with Error::Io(IoError::ModelStructureMismatch).

The position matters for a second reason. A checkpoint path is <scope>.<name>, and scope is the index of the layer counted from the input. A pooling layer that is present but absent from the file therefore shifts the index of every layer behind it. The paths of those layers then stop matching. The layer-count check catches that case first.

You cannot drop a MaxPooling2D from the rebuilt model to save space. The structural checkpoint depends on that layer being present.

use rustyml::neural_network::Shape;
use rustyml::neural_network::sequential::SequentialBuilder;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::optimizers::*;
use rustyml::neural_network::losses::*;
use ndarray::Array;

fn main() {
    let x = Array::from_elem((2, 4, 4, 3), 0.5f32).into_dyn();
    let y = Array::ones((2, 2)).into_dyn();

    let mut model = SequentialBuilder::new()
        .add(GlobalAveragePooling2D::new())
        .add(Dense::new(2, Activation::ReLU).unwrap())
        .build(&Shape::known(x.shape()))
        .unwrap();
    model.compile(RMSprop::new(0.001, 0.9, 1e-8, 0.0).unwrap(), MeanSquaredError::new());
    model.fit(&x, &y, 2).unwrap();
    model.save_to_path("pool_head.bin").unwrap();

    // Rebuild the same architecture, then load. The pooling layer carries no weights.
    // It must still occupy the same position, with the same type, for the structure check to pass.
    let mut restored = SequentialBuilder::new()
        .add(GlobalAveragePooling2D::new())
        .add(Dense::new(2, Activation::ReLU).unwrap())
        .build(&Shape::known(x.shape()))
        .unwrap();
    restored.load_from_path("pool_head.bin").unwrap();

    let out = restored.predict(&x).unwrap();
    assert_eq!(out.shape(), &[2, 2]);

    std::fs::remove_file("pool_head.bin").unwrap();
}

See 3.9. Saving and Loading Weights and 7.2. Model Persistence in Depth for the full serialization story. That story covers the postcard format, why save_to_path skips the optimizer and loss, and how the load checks weight shapes.

3.6.8. Performance and cost

The engine avoids extra allocations. It reads the whole input as 1 contiguous &[f32]. Its unit of work is a single output position. For each window tap, a serial loop folds channels adjacent floats into a channels-wide accumulator. The engine then assembles the output with 1 from_shape_vec call, instead of writing scalars one at a time.

Cost scales with output positions, times window volume, times channels. Overlapping windows (a stride smaller than the pool size) do proportionally more work than the non-overlapping default. Global pooling is the cheapest case. It runs 1 linear scan per batch item.

Forward work runs in parallel with rayon, split over (batch item, block of output positions). Each block writes a disjoint output slab. Splitting on positions, rather than on the batch alone, keeps every thread busy even when batch == 1.

Backward work splits differently, over (batch item, channel slab). A slab that owns channels [j0, j1) touches only input addresses in that range, modulo channels. So the scatter needs no halo, no merge step, and no atomic operations.

Both passes go parallel only after the estimated total work clears a threshold. That estimate is batch * output_positions * channels * window_volume element operations. The threshold, POOL_PARALLEL_MIN_OPS, defaults to 12,000, and you can override it through the tuning module. The gate counts total taps, not task count. This keeps the gate honest, whether the work sits in a few wide-channel positions or in many narrow ones. Small tensors run in serial, to avoid paying the rayon task overhead for no benefit.

See 7.3. Performance Tuning and Parallelism to move that gate and measure the effect.

1 warning carries over from the Layer contract. Backward is pure math, and it does not sanitize non-finite values. Max pooling also propagates NaN on purpose. So a NaN in a pooling input travels straight through, unchanged. Control large but finite gradients with optimizer-level global-norm clipping. Do not expect pooling to clamp them.

3.6.9. UpSampling: the inverse of pooling

Pooling shrinks the spatial axes. Upsampling grows them back. UpSampling1D, UpSampling2D, and UpSampling3D multiply the extent of every spatial axis by a whole number. Each layer leaves the batch axis and the channel axis untouched, and holds no parameter. This is the parameter-free way to write a decoder: an autoencoder, or a segmentation head that has to return to the input resolution. For a decoder that learns how it grows instead, see the transposed convolutions in 3.5. Convolutional Layers.

The 2D and 3D factor arguments take impl Into<Factor2D> and impl Into<Factor3D>. An integer gives the same factor to every spatial axis, and a tuple names 1 factor per axis. Read the tuple the way you read the border tuples in 3.5. Convolutional Layers. For example, UpSampling2D::new((2, 3), ...) doubles the height and triples the width. All 3 constructors return a Result, and a factor of 0 is Error::InvalidParameter, because it would empty the tensor.

The default mode repeats. The layer copies each input position into a block of factor output positions. A 2x upsampled image is therefore the original with every pixel grown into a 2x2 square. In shape that is the exact inverse of a (2, 2) pooling window. In value it is not, because pooling threw information away and no upsampling can invent it back.

use rustyml::neural_network::Shape;
use rustyml::neural_network::sequential::SequentialBuilder;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::optimizers::*;
use rustyml::neural_network::losses::*;
use ndarray::Array;

fn main() {
    // An encoder halves both spatial axes, and the decoder puts them back.
    let x = Array::from_shape_fn((2, 8, 8, 3), |(_, i, j, c)| (i + j + c) as f32 * 0.1).into_dyn();

    let mut model = SequentialBuilder::new()
        .add(MaxPooling2D::new((2, 2)))                             // -> [2, 4, 4, 3]
        .add(UpSampling2D::new(2, Interpolation::Nearest).unwrap()) // -> [2, 8, 8, 3]
        .build(&Shape::known(x.shape()))
        .unwrap();
    model.compile(SGD::new(0.01, 0.0, false, 0.0).unwrap(), MeanSquaredError::new());

    let out = model.predict(&x).unwrap();
    assert_eq!(out.shape(), &[2, 8, 8, 3]);

    // The shape comes back, but most values do not. Pooling kept 1 position out of every 4.
    let exact = out.iter().zip(x.iter()).filter(|&(a, b)| (a - b).abs() < 1e-6).count();
    assert!(exact < x.len());
}

Only UpSampling2D takes an interpolation mode. The 4 modes that are not the repeat resample with a separable kernel. A new pixel is therefore a weighted sum of its neighbors along each axis:

Every mode puts an output position at the center of the source region it covers, using the half-pixel convention. Output position j reads the source coordinate (j + 0.5) / factor - 0.5. Near an edge, part of the kernel falls outside the image. The layer drops those weights, and divides the remaining weights by their own sum. The weights of any output position therefore always sum to 1. A constant image stays constant, edges included.

use ndarray::Array;
use rustyml::neural_network::Ctx;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::UnaryLayer;

fn main() {
    // A 2x2 image, enlarged 2x with a triangle kernel.
    let x = Array::from_shape_vec((1, 2, 2, 1), vec![1.0, 2.0, 3.0, 4.0]).unwrap().into_dyn();
    let mut ctx = Ctx::inference();
    let mut smooth = UpSampling2D::new(2, Interpolation::Bilinear).unwrap();
    let out = smooth.forward_mut(&x, &mut ctx).unwrap();

    // A corner keeps its source value, because the kernel there is clipped to 1 position.
    assert!((out[[0, 0, 0, 0]] - 1.0).abs() < 1e-6);
    assert!((out[[0, 3, 3, 0]] - 4.0).abs() < 1e-6);
    // An interior position blends 4 neighbors: 0.5625*1 + 0.1875*2 + 0.1875*3 + 0.0625*4.
    assert!((out[[0, 1, 1, 0]] - 1.75).abs() < 1e-6);

    // The weights of an output position add up to 1, so a constant image stays constant.
    let flat = Array::from_elem((1, 3, 4, 2), 7.5f32).into_dyn();
    for mode in [Interpolation::Bicubic, Interpolation::Lanczos3, Interpolation::Lanczos5] {
        let mut layer = UpSampling2D::new(3, mode).unwrap();
        let same = layer.forward_mut(&flat, &mut ctx).unwrap();
        assert!(same.iter().all(|v| (v - 7.5).abs() < 1e-5));
    }
}

The cubic and Lanczos kernels have negative lobes, so their output can leave the input range. A bright pixel on a dark field shows a dark halo, and a map of probabilities can go slightly below 0. When the range matters, clamp after the layer. The repeat mode and bilinear cannot overshoot, because all their weights are positive.

The backward pass is the transpose of that same weighted gather. Every input position collects the gradient of each output position it fed, under the same weight. For the repeat mode this is a plain sum: a position that fed factor output positions gets the total of their factor gradients.

The Lanczos kernels have large lobes that cancel, so computing the weight table in f32 loses precision before it touches a pixel. This layer builds the table in f64 and rounds once at the end. Over 250 shape, factor, and mode combinations, the worst deviation of Lanczos5 against an independent f64 reference is 2.5e-7.

The layer pays 1 multiply-add per position read per output element, once per spatial axis rather than once per output pixel. The repeat mode copies whole runs of channels at memory speed, while Lanczos5 reads 11 positions and costs about 11 times as much. Each axis pass goes parallel with rayon once output elements x positions read clears tuning::upsampling::set_parallel_min_ops, which defaults to 2,000,000. The gate counts element ops rather than elements, because the repeat mode reads 1 position and Lanczos5 reads 11. An element count alone would put 2 passes of very different cost on the same side of the gate. The layer adds the positions in a fixed order, so the parallel path returns the same bits as the serial path.