Align compatible shapes, reduce a named axis

Predict one feature offset across two token rows

Modern decoder code frequently carries activations with a final feature axis. Begin with two token positions, each holding three feature values:

tokens shape [2,3]
[[1, 2, 3],
 [4, 5, 6]]

The feature bias has no token axis:

bias shape [3]
[10, 20, 30]

Before reading the reveal, predict a result shape and six sums. Should the bias be rejected because it has rank one, copied into a new [2,3] input first, or mapped to each token row by a fixed compatibility rule?

Align the shapes from their trailing ends:

tokens: [2,3]
bias:   [1,3]   # missing leading dimension acts as 1
result: [2,3]

Feature extent $3$ matches $3$. On the leading aligned axis, size $1$ can be reused across size $2$. Output coordinate $(r,c)$ therefore reads token coordinate $(r,c)$ and bias coordinate $(c)$. The result is:

Now predict before calculating: sum across token axis 0, mean across feature axis 1 while retaining it, and max across feature axis 1 while removing it. The results are respectively shape [3] with [25,47,69], shape [2,1] with [22,25], and shape [2] with [33,36].

Map broadcast coordinates and reduce one axis

Broadcasting and mean reduction can be written together:

The first expression keeps tensor shape rules separate from scalar arithmetic. $y_{\mathbf{i}}$ is one output value at complete coordinate $\mathbf{i}$. Function $f$ can add, multiply, or perform any other scalar calculation. It receives the left value from tensor $a$ at $\beta_a(\mathbf{i})$ and the right value from tensor $b$ at $\beta_b(\mathbf{i})$.

Each beta mapping removes output axes that were missing from that input and selects coordinate zero on an aligned input axis of size one. Other coordinates pass through. In the frozen example:

The second expression fixes the coordinates $\mathbf{i}_{-k}$ outside axis $k$, lets $i_k$ visit the $n_k$ entries on that axis, and divides their ordered sum by $n_k$. The mean is defined only when $n_k>0$; the implementation returns a typed empty-axis error otherwise. The complete coordinate $\mathbf{i}$ combines those fixed coordinates with the current $i_k$. With keep_dim=true, the selected axis remains in its original output position with extent one. This does not add a value to the mean or change its divisor.

Account for every coordinate, extent, and mapping

Compatibility is checked axis by axis after right alignment. Two extents are compatible when they are equal or at least one is one. Missing leading extents behave as one. The output takes the equal extent, or otherwise the non-one extent. This last wording matters: zero with one produces zero, not one.

A scalar has shape [] and broadcasts to any result because it has no axes to contradict. A scalar cannot be reduced by this operation because it has no valid axis. An empty output enumerates no coordinates, so the supplied scalar function is never called.

From fixed context to tensor-wide decoder math

Bengio et al. describe n-gram models as conditional-probability tables for a fixed number of preceding words; their neural language model concatenates learned context-word features, uses a hyperbolic-tangent hidden layer, and produces next-word probabilities with softmax. Its prediction is still organized around one selected fixed window rather than every position’s available causal prefix and the explicit batch, sequence, and head axes used by later decoder Transformers.

For an $n$-gram of order $n$, that fixed context contains $n-1$ preceding words.

The earlier source is Bengio et al., A Neural Probabilistic Language Model.

That model improves on separate count rows by learning distributed word features and shared neural parameters, but its input is still a selected fixed window. The elementwise $\tanh$ and vocabulary softmax are genuine model computations; this chapter does not imply that the paper prescribed a general tensor interface.

The later sources are Vaswani et al., Attention Is All You Need and OpenAI’s official GPT-2 model.py. Vaswani et al. define masked decoder self-attention over query, key, and value matrices, apply softmax to scaled query-key scores, wrap each sublayer with a residual connection followed by layer normalization, and apply the same feed-forward network separately and identically at every position. The official GPT-2 implementation labels batch, sequence, feature, head, destination, and source axes. Its softmax subtracts a maximum computed over the last axis, exponentiates the shifted values, and divides by their sum over that axis; both reductions retain the axis. Its normalization takes last-axis means before applying feature-sized scale and bias vectors.

Those later calculations keep several meanings visible at once: batch item, token position, attention head, destination position, source position, and feature. Elementwise activation, residual addition, masks, scales, and affine feature parameters must apply across those axes. Softmax and normalization must reduce the intended axis while retaining enough shape to combine the result with the original tensor.

Broadcasting and explicit-axis reductions let this course apply scalars or feature-sized parameters across decoder tensors and compute the per-axis statistics needed by attention softmax and feature normalization. The exact trailing-axis rule, shape errors, empty-axis behavior, keep-dimension option, and allocation policy belong to this implementation; the model sources specify the computations, while the NumPy guide documents the supporting shape-alignment rule.

The tiny Rust contrast keeps that boundary visible. fixed_context_feature_step performs one shape-specific three-feature calculation. The general path applies the same feature offset over two token rows and derives named reductions. It is a calculation-sized bridge from one context representation to explicit token-feature axes, not a Transformer innovation or a complete implementation of either cited model.

/// Applies one fixed-width feature offset to one context representation.
///
/// This shape-specific baseline stands for an earlier one-example calculation;
/// it is not presented as the exact equation of any cited model.
pub fn fixed_context_feature_step(context: [f64; 3], feature_bias: [f64; 3]) -> [f64; 3] {
    [
        context[0] + feature_bias[0],
        context[1] + feature_bias[1],
        context[2] + feature_bias[2],
    ]
}

/// Runs the same scalar operations over explicit token and feature axes.
pub fn tiny_token_feature_example() -> Result<TinyTokenFeatureExample, TensorOpError> {
    let tokens = Tensor::from_vec(TOKEN_SHAPE.to_vec(), TOKEN_VALUES.to_vec())?;
    let bias = Tensor::from_vec(BIAS_SHAPE.to_vec(), BIAS_VALUES.to_vec())?;
    let biased = map_binary(&tokens.view(), &bias.view(), |value, offset| value + offset)?;
    let squared = map_unary(&tokens.view(), |value| value * value)?;
    let sum_axis_0 = sum_axis(&biased.view(), 0, false)?;
    let mean_axis_1_kept = mean_axis(&biased.view(), 1, true)?;
    let max_axis_1 = max_axis(&biased.view(), 1, false)?;

    Ok(TinyTokenFeatureExample {
        tokens,
        bias,
        biased,
        squared,
        sum_axis_0,
        mean_axis_1_kept,
        max_axis_1,
    })
}

Plan shapes before evaluating values

broadcast_shape is pure shape planning. It right-aligns ranks, treats missing leading dimensions as one, and visits aligned output axes from left to right. Equal extents pass through. If one extent is one, the other wins. Otherwise the error records the exact output axis and both sizes.

After compatibility succeeds, the planner asks the existing tensor layout code to validate all suffix products. Thus [0,usize::MAX,1] with [1,1,2] reaches ShapeOverflow before allocation: the outer zero makes the input valid, but the planned suffix multiplies the usize::MAX extent by $2$, which does not fit. Compatibility errors take precedence over this layout check.

/// Computes the checked output shape for trailing-axis broadcasting.
///
/// Missing leading dimensions act as size one. Aligned dimensions are
/// compatible when they are equal or either one is size one. Compatibility is
/// reported from the leftmost aligned output axis before layout overflow.
pub fn broadcast_shape(left: &[usize], right: &[usize]) -> Result<Vec<usize>, TensorOpError> {
    let output_rank = left.len().max(right.len());
    let left_padding = output_rank - left.len();
    let right_padding = output_rank - right.len();
    let mut output = Vec::with_capacity(output_rank);

    for axis in 0..output_rank {
        let left_dimension = left
            .get(axis.wrapping_sub(left_padding))
            .copied()
            .unwrap_or(1);
        let right_dimension = right
            .get(axis.wrapping_sub(right_padding))
            .copied()
            .unwrap_or(1);
        let dimension = if left_dimension == right_dimension {
            left_dimension
        } else if left_dimension == 1 {
            right_dimension
        } else if right_dimension == 1 {
            left_dimension
        } else {
            return Err(TensorOpError::IncompatibleBroadcast {
                axis,
                left_dimension,
                right_dimension,
            });
        };
        output.push(dimension);
    }

    checked_row_major_layout(&output)?;
    Ok(output)
}

The compatibility rule matches the NumPy broadcasting guide on the supporting rule that trailing dimensions are compatible when they are equal or one is size one. This implementation additionally makes the zero-with-one case explicit: the non-one extent is zero, so the result remains empty. The guide is not a source for the LLM history in this chapter and does not prescribe this implementation’s errors, logical-view traversal, or owned-output policy.

TensorView::get remains the public path for one coordinate supplied by a caller. That coordinate may have the wrong rank or an out-of-bounds axis, so every call validates it before reading storage. Elementwise kernels receive different inputs: their views are already valid, and map_binary validates the equal-or-one broadcast shape before value traversal begins. Each internal stride plan is then checked once and reused.

For map_binary, each input receives one effective element stride for every output axis. A missing leading input axis or an aligned input axis of extent one receives effective stride 0, because changing that output coordinate must select the same input value. Every other aligned axis keeps its TensorView stride. In the frozen example, token strides [3,1] yield source offsets [0,1,2,3,4,5], while bias effective strides [0,1] yield source offsets [0,1,2,0,1,2].

Pairing those two offset sequences implements the six conceptual coordinate mappings predicted earlier. map_unary uses the same checked offset traversal with the input view’s own strides. Neither scalar loop constructs a coordinate vector or re-enters the public coordinate validator. The change affects only traversal: it preserves left-before-right closure arguments, logical row-major closure order, ordinary safe bounds-checked storage reads, and allocation of a new contiguous result. Empty outputs call the closure zero times. A transposed or sliced view supplies different strides without an implicit materialize.

All output reservations are fallible. For example, [usize::MAX,0] is a valid empty input shape, but summing its empty axis would require an owned result with usize::MAX zeros. The implementation returns OutputAllocationFailed { elements: usize::MAX } before iteration instead of letting a capacity request panic.

/// Applies one scalar function in logical row-major order and owns the result.
pub fn map_unary<F>(input: &TensorView<'_>, mut operation: F) -> Result<Tensor, TensorOpError>
where
    F: FnMut(f64) -> f64,
{
    let mut values = output_buffer(input.len())?;
    for input_offset in input.logical_offsets() {
        values.push(operation(input.value_at_storage_offset(input_offset)));
    }
    Tensor::from_vec(input.shape().to_vec(), values).map_err(Into::into)
}

/// Applies one scalar function across two trailing-axis-compatible views.
pub fn map_binary<F>(
    left: &TensorView<'_>,
    right: &TensorView<'_>,
    mut operation: F,
) -> Result<Tensor, TensorOpError>
where
    F: FnMut(f64, f64) -> f64,
{
    let output_shape = broadcast_shape(left.shape(), right.shape())?;
    let (_, output_len) = checked_row_major_layout(&output_shape)?;
    let mut values = output_buffer(output_len)?;

    let left_strides = broadcast_effective_strides(left, output_shape.len());
    let right_strides = broadcast_effective_strides(right, output_shape.len());
    let left_offsets = left
        .projected_offsets(&output_shape, &left_strides, output_len)
        .expect("a compatible broadcast retains a valid left traversal plan");
    let right_offsets = right
        .projected_offsets(&output_shape, &right_strides, output_len)
        .expect("a compatible broadcast retains a valid right traversal plan");

    for (left_offset, right_offset) in left_offsets.zip(right_offsets) {
        let left_value = left.value_at_storage_offset(left_offset);
        let right_value = right.value_at_storage_offset(right_offset);
        values.push(operation(left_value, right_value));
    }

    Tensor::from_vec(output_shape, values).map_err(Into::into)
}

The three reductions share axis validation, output-shape construction, and fixed ascending traversal. Sum over an empty selected axis has the additive identity 0.0. Mean and max have no value there and return distinct typed errors. A zero extent on another retained axis simply produces an empty output.

After the selected axis and output layout have been validated, each nonempty output group receives one source base offset. The reduction visits that group by repeatedly adding the selected input axis’s element stride, in ascending axis-coordinate order. For the contiguous [2,3] result, axis 0 uses bases [0,1,2] and stride 3, producing source-offset groups [0,3], [1,4], and [2,5]. Axis 1 uses bases [0,3] and stride 1, producing [0,1,2] and [3,4,5]. A transposed or sliced view supplies different bases and strides, but the logical groups and accumulation order remain unchanged. An empty selected sum writes 0.0 without reading a base offset. If a different, non-selected axis has extent zero, the output has zero groups and no source read occurs.

For each nonempty group, maximum starts with the value at selected-axis coordinate zero, replaces it only for a strictly larger value, and retains the first NaN it encounters. Equal values retain the earlier exact bits, including -0.0 versus +0.0. This explicit policy avoids inheriting accidental NaN or tie behavior from a convenience function.

/// Sums one explicit axis in ascending index order.
///
/// An empty selected axis uses the additive identity, so every output group is
/// `0.0`. `keep_dim` replaces the selected extent with one instead of removing
/// the axis.
pub fn sum_axis(
    input: &TensorView<'_>,
    axis: usize,
    keep_dim: bool,
) -> Result<Tensor, TensorOpError> {
    reduce_axis(input, axis, keep_dim, Reduction::Sum)
}

/// Averages one explicit nonempty axis in ascending index order.
pub fn mean_axis(
    input: &TensorView<'_>,
    axis: usize,
    keep_dim: bool,
) -> Result<Tensor, TensorOpError> {
    reduce_axis(input, axis, keep_dim, Reduction::Mean)
}

/// Selects the maximum over one explicit nonempty axis.
///
/// The fold propagates the first NaN and keeps the earlier value on equal
/// comparisons, including the earlier signed-zero bit pattern.
pub fn max_axis(
    input: &TensorView<'_>,
    axis: usize,
    keep_dim: bool,
) -> Result<Tensor, TensorOpError> {
    reduce_axis(input, axis, keep_dim, Reduction::Max)
}

#[derive(Clone, Copy)]
enum Reduction {
    Sum,
    Mean,
    Max,
}

fn reduce_axis(
    input: &TensorView<'_>,
    axis: usize,
    keep_dim: bool,
    reduction: Reduction,
) -> Result<Tensor, TensorOpError> {
    if axis >= input.rank() {
        return Err(TensorOpError::ReductionAxisOutOfBounds {
            axis,
            rank: input.rank(),
        });
    }

    let axis_len = input.shape()[axis];
    match reduction {
        Reduction::Mean if axis_len == 0 => {
            return Err(TensorOpError::EmptyMeanAxis { axis });
        }
        Reduction::Max if axis_len == 0 => {
            return Err(TensorOpError::EmptyMaxAxis { axis });
        }
        _ => {}
    }

    let output_shape = reduction_shape(input.shape(), axis, keep_dim);
    let (_, output_len) = checked_row_major_layout(&output_shape)?;
    let mut values = output_buffer(output_len)?;

    if axis_len == 0 {
        debug_assert!(matches!(reduction, Reduction::Sum));
        values.resize(output_len, 0.0);
        return Tensor::from_vec(output_shape, values).map_err(Into::into);
    }

    let group_strides = reduction_group_strides(input.strides(), axis, keep_dim);
    let group_offsets = input
        .projected_offsets(&output_shape, &group_strides, output_len)
        .expect("a checked reduction retains a valid group traversal plan");
    let axis_stride = input.strides()[axis];

    for group_offset in group_offsets {
        let value = match reduction {
            Reduction::Sum | Reduction::Mean => {
                let mut total = 0.0;
                let mut input_offset = group_offset;
                for index in 0..axis_len {
                    total += input.value_at_storage_offset(input_offset);
                    if index + 1 < axis_len {
                        input_offset = input_offset
                            .checked_add(axis_stride)
                            .expect("a checked view cannot overflow along a reduction axis");
                    }
                }
                if matches!(reduction, Reduction::Mean) {
                    total / axis_len as f64
                } else {
                    total
                }
            }
            Reduction::Max => {
                let mut input_offset = group_offset;
                let mut maximum = input.value_at_storage_offset(input_offset);
                for _ in 1..axis_len {
                    input_offset = input_offset
                        .checked_add(axis_stride)
                        .expect("a checked view cannot overflow along a reduction axis");
                    let candidate = input.value_at_storage_offset(input_offset);
                    if !maximum.is_nan() && (candidate.is_nan() || candidate > maximum) {
                        maximum = candidate;
                    }
                }
                maximum
            }
        };
        values.push(value);
    }

    Tensor::from_vec(output_shape, values).map_err(Into::into)
}

The learner program exercises the frozen model-shaped case, scalar broadcast, an empty broadcast whose closure must run zero times, the empty-sum identity, and each typed error. The integer-valued calculations can be compared exactly. For the mean of [0.1,0.2,0.3], an absolute tolerance of 1e-12 is appropriate because those decimals are not exact binary floating-point values.

    let example = tiny_token_feature_example()?;

    let scalar = Tensor::from_vec(vec![], vec![0.5])?;
    let scalar_broadcast = map_binary(&example.tokens.view(), &scalar.view(), |value, offset| {
        value + offset
    })?;

    let empty = empty_fixture()?;
    let mut closure_calls = 0;
    let empty_broadcast = map_binary(&empty.view(), &example.bias.view(), |value, offset| {
        closure_calls += 1;
        value + offset
    })?;
    let empty_sum = sum_axis(&empty.view(), 1, false)?;

    let incompatible = Tensor::from_vec(vec![2], vec![1.0, 2.0])?;
    let broadcast_error = map_binary(
        &example.tokens.view(),
        &incompatible.view(),
        |left, right| left + right,
    )
    .unwrap_err();
    let mean_error = mean_axis(&empty.view(), 1, false).unwrap_err();
    let max_error = max_axis(&empty.view(), 1, false).unwrap_err();
    let scalar_reduction_error = sum_axis(&scalar.view(), 0, false).unwrap_err();

The example keeps the chapter boundary explicit: it demonstrates scalar maps, shape alignment, named reductions, and typed failures, but it does not implement softmax, normalization, or matrix multiplication. Those later operations will reuse these checked primitives.

See reused features and reductions along named axes

Read the figure in four passes:

1. Compare each original shape with its trailing-axis alignment.
2. Follow all six output coordinates to one token coordinate and one bias coordinate. Notice that each bias coordinate appears in both token rows.
3. Compare the flat groups consumed by axis-0 sum and axis-1 mean/max. Check the keep dimension field before reading each output shape.
4. Inspect the three rejected records: incompatible sizes 3 and 2, empty mean, and empty max.

The reuse marker identifies a bias coordinate selected for both token rows; it does not mean the bias tensor was copied first. Downward markers introduce the groups combined by each reduction. Rejection markers pair every failed request with the incompatible extents or empty selected axis that makes a result undefined.

Predict valid shapes and reduction results

Write the shape and coordinate mapping before calculating values.

1. Align [2,1,3] with [4,3]. Predict the output shape and both source coordinates for output [1,2,0].
2. Decide whether [2,3] and [2] broadcast. If not, name the output axis and the two incompatible extents.
3. Broadcast [0,3] with [1,3]. Predict the result shape and closure call count.
4. For the biased example, predict axis-0 sum, axis-1 mean with the dimension retained, and axis-1 max without it.
5. Reduce selected empty axis 1 of shape [2,0,3] with sum, mean, and max.
6. Mean-reduce rank-one values [0.1,0.2,0.3] on axis 0 without retaining it. Predict the output shape and explain the tolerance.
7. Apply the maximum policy to [1,NaN_A,NaN_B] and [-0.0,+0.0,-1.0]. Predict the surviving payload and zero sign.
8. For token strides [3,1] and bias stride [1], write both effective broadcast stride plans and their six source-offset sequences. Then write the axis-0 reduction bases, selected-axis stride, and source-offset groups. Why can these internal plans be reused while TensorView::get must validate every caller-supplied coordinate?

Misconception check: broadcasting is not an eager command to copy a smaller tensor until shapes look equal. It is a coordinate rule for compatible axes; this implementation evaluates directly through those mappings and owns only the result. A reduction is also not “sum everything”: it must name the axis whose coordinate changes while every other coordinate stays fixed.

Prepare the primitives behind normalization and softmax

The cumulative tensor core can now combine scalar or feature-sized values with token tensors, apply unary activations to logical views, and compute fixed-order sum, mean, or max over one explicit axis. The results are checked owned tensors, so later model code has one predictable shape and storage boundary.

Public coordinate lookup still validates every coordinate supplied by a caller. Inside a validated operation, zero effective broadcast strides and one checked base-plus-stride plan per reduction group express the same coordinate rules without rebuilding and revalidating a coordinate for every scalar.

These are supporting pieces, not finished model layers. Stable attention softmax will subtract an axis maximum, exponentiate elementwise, sum the same axis with its dimension retained, and divide. Feature normalization will compute feature-axis statistics and combine the normalized values with feature-sized parameters. Those complete algorithms will be developed in later chapters.

Chapter 11 first adds matrix multiplication. Unlike broadcasting, its contracted dimensions must match as a multiplication invariant, and its output dimensions come from rows and columns rather than the equal-or-one rule. Keeping that boundary explicit prevents a convenient elementwise planner from silently standing in for learned linear transformations.