12 · Content revision 8

Turn extreme logits into stable probabilities

Turn vocabulary and attention logits into stable probabilities, log-probabilities, log-sum-exp values, and indexed mean NLL with dependency-free Rust.

Predict three shifted rows

Start with three rows of two class scores:

shape [3,2]
[[    0,     1],
 [ 1000,  1001],
 [-1001, -1000]]

These values are logits: scores before normalization, not probabilities. The second entry is larger by exactly one in every row. Predict the result of subtracting each row’s maximum before reaching for a calculator:

\begin{aligned} [0,1]-1 &= [-1,0], \\ [1000,1001]-1001 &= [-1,0], \\ [-1001,-1000]-(-1000) &= [-1,0]. \end{aligned}

Exponentiating now produces [0.367879441171,1] for every row. Dividing by their sum 1.367879441171 gives the same pair every time:

[0.268941421370, 0.731058578630]

That equality is not an approximation caused by the display: the Rust example calculates each shifted row identically. The maximum shift removes a shared constant while preserving the one-unit difference that controls their relative probability.

For targets [1,0,1], predict which entry each row contributes. Their negative log-probabilities are [0.313261687518,1.313261687518,0.313261687518], so the mean is 0.646595020852 nats per target.

Normalize with one maximum shift

The chapter’s stable softmax formula is:

p_i=\frac{\exp(\ell_i-m)}{\sum_j\exp(\ell_j-m)}, \quad m=\max_j\ell_j

The largest shifted logit is exactly zero, so at least one exponential is one and none is larger than one. This removes the avoidable raw overflow in $\exp(1000)$ and prevents an all-zero denominator for the very negative row.

In exact arithmetic, adding the same constant to every logit also adds it to the maximum, so every difference stays unchanged. The worked rows preserve those differences exactly in f64; an arbitrary floating-point shift can instead round away a small difference before softmax sees it. The maximum is not a probability and subtraction alone does not normalize anything; exponentiation and division by the complete group sum still matter.

The same shifted terms support three log-domain quantities:

\operatorname{LSE}(\ell_1,\ldots,\ell_n) =m+\ln\sum_j\exp(\ell_j-m)

\log p_i=(\ell_i-m)-\ln\sum_j\exp(\ell_j-m)

\mathcal{L}_t=(m-\ell_t)+\ln\sum_j\exp(\ell_j-m)

Log-sum-exp adds the maximum back after taking the natural logarithm of the shifted sum. Log-softmax keeps the safer shifted difference and subtracts that logarithm. Indexed NLL selects one target and evaluates its loss without first rounding through an ordinary probability.

For one normalization group, define the shared shifted-exponential sum:

S=\sum_j\exp(\ell_j-m)

The maximum $m$ , shifted-exponential sum $S$ , and log-normalizer $\ln S$ are shared by every class in that group. Softmax uses $m$ and $S$ ; log-softmax uses $m$ and $\ln S$ ; indexed NLL uses $m$ , the selected target logit $\ell_t$ , and $\ln S$ . A training operation that also needs probabilities for its future backward gradient calculation can emit them from the same three group facts instead of calculating the group statistics again.

Name each probability quantity

Symbol	Operational meaning
$p_i$	The normalized probability assigned to class $i$ .
$\ell_i$	The finite input logit for class $i$ .
$m$	The largest logit in the selected normalization group.
$i$	The class whose probability is being computed.
$j$	The class index traversed across the complete denominator.

An axis divides a tensor into independent normalization groups. For shape [3,2] and axis 1, there are three groups with two classes each. Removing the class axis gives group shape [3], so flat targets [1,0,1] correspond to rows zero, one, and two in that order.

Log-sum-exp may remove the selected axis or retain it with extent one. Softmax and log-softmax preserve the input shape. Every successful tensor result owns a contiguous row-major buffer even when the input is a sliced or transposed view.

From vocabulary softmax to Transformer probabilities

Bengio et al.’s neural language model uses an output softmax to turn vocabulary scores into positive next-word probabilities that sum to one. In finite precision, directly exponentiating unshifted large logits can overflow, while directly exponentiating sufficiently negative logits can round every term to zero.

The earlier checkpoint is Bengio et al., A Neural Probabilistic Language Model. Bengio et al. describe an output softmax whose values are positive and sum to one, interpreting its inputs as unnormalized log probabilities for the next word.

That output distribution is the relevant historical step: learned neural scores become probabilities over a vocabulary. The cited source does not say that its implementation used this chapter’s literal Rust baseline, arbitrary-axis API, finite-input rejection, or error order.

The Transformer reuses softmax for scaled query-key scores inside attention and for next-token predictions. OpenAI’s published GPT-2 source shows a stable implementation for attention: subtract the maximum along the last axis before exponentiating, sum the shifted exponentials, and normalize before combining values.

The later sources are Vaswani et al., Attention Is All You Need and OpenAI’s published GPT-2 model.py. Vaswani et al. define scaled dot-product attention by applying softmax to scaled query-key products before weighting values, and apply a learned linear transform plus softmax to decoder outputs for predicted next-token probabilities. OpenAI’s GPT-2 source implements last-axis softmax by subtracting the maximum with retained dimensions, exponentiating, and dividing by the sum with retained dimensions; its attention path applies that helper to scaled masked scores before combining values. The source names those reductions reduce_max and reduce_sum.

In exact arithmetic, adding one constant to every logit leaves softmax unchanged. Maximum shifting preserves that distribution while avoiding raw-exponential failures for the worked rows. Log-sum-exp supplies the stable log-normalizer; log-softmax retains class scores in the log domain, and fused indexed mean NLL retains a target loss when the corresponding ordinary probability rounds to zero. This course’s arbitrary-axis API, finite-input policy, target layout, allocation rules, and error precedence are local correctness decisions.

The bounded Rust contrast starts with a literal formula translation. It succeeds for [0,1]; raw positive extremes become infinity divided by infinity, and raw negative extremes become zero divided by zero. This demonstrates the numerical pressure that stable normalization solves on the path to modern LLMs. This is floating-point behavior rather than a programming-language rule: implementations face the same overflow and underflow pressure even when their APIs and error handling differ.

Expose raw-exponential normalization for one ordinary and two extreme finite rows rust/demos/ch12-stable-softmax/src/lib.rs#direct-output-softmax

/// Applies the literal exponential normalization used as a bounded baseline.
///
/// This exposes finite-precision overflow and underflow; it is not attributed
/// to the software implementation of any cited language model.
pub fn direct_output_softmax(logits: &[f64]) -> Vec<f64> {
    let exponentials = logits.iter().map(|value| value.exp()).collect::<Vec<_>>();
    let denominator = exponentials.iter().sum::<f64>();
    exponentials
        .into_iter()
        .map(|value| value / denominator)
        .collect()
}

Implement checked log-domain operations

The error surface keeps rejected invariants visible. Axis bounds come first. Softmax, log-softmax, and indexed NLL then reject an empty class axis. Tensor-returning operations check their output layout and reserve output storage before reading logits. Indexed NLL instead checks group layout, target count, a nonempty mean, and every target bound before reading any logit. Only then can the first NaN or signed infinity be reported in group-major, class-minor order.

Keep axes, empty classes, outputs, non-finite logits, and indexed targets distinct rust/crates/llm-from-scratch/src/nn/probability.rs#probability-errors

/// A rejected probability operation, target, output, or converted view operation.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ProbabilityError {
    /// An owned output layout violates the tensor storage invariant.
    Tensor(TensorError),
    /// A tensor-view error was converted into the probability error type.
    View(TensorViewError),
    /// The requested class axis does not exist.
    AxisOutOfBounds { axis: usize, rank: usize },
    /// Softmax, log-softmax, and indexed NLL need at least one class.
    EmptyNormalizationAxis { axis: usize },
    /// The checked output shape is valid, but its value buffer cannot be reserved.
    OutputAllocationFailed { elements: usize },
    /// The first rejected logit in group-major, class-minor order is NaN.
    NaNLogit { group: usize, class: usize },
    /// The first rejected logit in group-major, class-minor order is positive infinity.
    PositiveInfinityLogit { group: usize, class: usize },
    /// The first rejected logit in group-major, class-minor order is negative infinity.
    NegativeInfinityLogit { group: usize, class: usize },
    /// There must be one flat target for every class-axis group.
    TargetCountMismatch { expected: usize, actual: usize },
    /// A mean is undefined when there are no target groups.
    EmptyTargets,
    /// One target does not name a class on the selected axis.
    TargetOutOfBounds {
        group: usize,
        target: usize,
        classes: usize,
    },
}

impl fmt::Display for ProbabilityError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Tensor(error) => error.fmt(formatter),
            Self::View(error) => error.fmt(formatter),
            Self::AxisOutOfBounds { axis, rank } => {
                write!(
                    formatter,
                    "probability axis {axis} is out of bounds for rank {rank}"
                )
            }
            Self::EmptyNormalizationAxis { axis } => {
                write!(formatter, "probability axis {axis} has no classes")
            }
            Self::OutputAllocationFailed { elements } => write!(
                formatter,
                "cannot allocate probability output for {elements} f64 values"
            ),
            Self::NaNLogit { group, class } => {
                write!(formatter, "logit at group {group}, class {class} is NaN")
            }
            Self::PositiveInfinityLogit { group, class } => write!(
                formatter,
                "logit at group {group}, class {class} is positive infinity"
            ),
            Self::NegativeInfinityLogit { group, class } => write!(
                formatter,
                "logit at group {group}, class {class} is negative infinity"
            ),
            Self::TargetCountMismatch { expected, actual } => write!(
                formatter,
                "indexed mean NLL needs {expected} targets, but received {actual}"
            ),
            Self::EmptyTargets => formatter.write_str("indexed mean NLL needs at least one target"),
            Self::TargetOutOfBounds {
                group,
                target,
                classes,
            } => write!(
                formatter,
                "target {target} at group {group} is out of bounds for {classes} classes"
            ),
        }
    }
}

impl Error for ProbabilityError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::Tensor(error) => Some(error),
            Self::View(error) => Some(error),
            _ => None,
        }
    }
}

impl From<TensorError> for ProbabilityError {
    fn from(error: TensorError) -> Self {
        Self::Tensor(error)
    }
}

impl From<TensorViewError> for ProbabilityError {
    fn from(error: TensorViewError) -> Self {
        Self::View(error)
    }
}

Before reading a logit, AxisPlan removes the selected class axis from the input shape and input strides. The remaining shape and strides define a checked cursor. Each value emitted by that cursor is the zero-based offset of class zero in one normalization group, measured in f64 elements from the start of the tensor owner’s flat storage. This chapter calls that value the group-base offset. The stride on the removed axis becomes the class stride. Starting from one group base, adding the class stride once selects class $1$ , adding it twice selects class $2$ , and so on.

For the contiguous worked tensor, shape [3,2] has source strides [2,1] and class axis 1. Removing that axis leaves group shape [3] and group stride [2], so the group-base cursor emits offsets [0,2,4]. The class stride is 1. The three groups therefore read source offsets [0,1], [2,3], and [4,5]. Those offsets correspond to rows zero, one, and two in the same order as flat targets [1,0,1].

For each nonempty normalization group, the first scan finds the maximum $m$ . The second scan resets to the same group base, visits classes in ascending order, separates one $\exp(0)=1$ term, and accumulates the remaining shifted exponentials as tail. It then records both $S=1+\mathrm{tail}$ for probability division and $\ln S=\ln(1+\mathrm{tail})$ , evaluated with ln_1p, for log-domain results. If several classes tie for the maximum, only one unit term is separated; the other tied classes remain in tail. Every class in the group uses these same three facts. ln_1p preserves a representable subnormal log-domain correction that ordinary $\ln(1+\mathrm{tail})$ could round away before the logarithm.

Compute one reusable statistics bundle for each checked normalization group rust/crates/llm-from-scratch/src/nn/probability.rs#checked-probability-groups

#[derive(Debug)]
struct AxisPlan {
    axis: usize,
    classes: usize,
    group_shape: Vec<usize>,
    group_strides: Vec<usize>,
    groups: usize,
    class_stride: usize,
}

#[derive(Clone, Copy, Debug)]
struct RowStats {
    maximum: f64,
    shifted_exponential_sum: f64,
    log_shifted_exponential_sum: f64,
}

#[derive(Clone, Copy, Debug)]
struct FiniteLogits;

#[derive(Clone, Copy, Debug)]
enum LogitFiniteness {
    Check,
    Validated(FiniteLogits),
}

impl AxisPlan {
    fn new(
        input: &TensorView<'_>,
        axis: usize,
        allow_empty_axis: bool,
    ) -> Result<Self, ProbabilityError> {
        if axis >= input.rank() {
            return Err(ProbabilityError::AxisOutOfBounds {
                axis,
                rank: input.rank(),
            });
        }

        let classes = input.shape()[axis];
        if classes == 0 && !allow_empty_axis {
            return Err(ProbabilityError::EmptyNormalizationAxis { axis });
        }

        let mut group_shape = input.shape().to_vec();
        group_shape.remove(axis);
        let mut group_strides = input.strides().to_vec();
        let class_stride = group_strides.remove(axis);
        let (_, groups) = checked_row_major_layout(&group_shape)?;
        Ok(Self {
            axis,
            classes,
            group_shape,
            group_strides,
            groups,
            class_stride,
        })
    }

    fn group_offsets(&self, input: &TensorView<'_>) -> StridedOffsets {
        input
            .projected_offsets(&self.group_shape, &self.group_strides, self.groups)
            .expect("a checked probability plan retains valid group-base offsets")
    }

    fn output_group_offsets(&self, output_strides: &[usize], output_len: usize) -> StridedOffsets {
        let mut group_strides = output_strides.to_vec();
        group_strides.remove(self.axis);
        StridedOffsets::checked(
            &self.group_shape,
            &group_strides,
            0,
            self.groups,
            output_len,
        )
        .expect("a checked probability output retains valid group-base offsets")
    }

    fn target_offset(&self, group_base: usize, target: usize) -> usize {
        let class_offset = target
            .checked_mul(self.class_stride)
            .expect("a checked probability plan cannot overflow a class offset");
        group_base
            .checked_add(class_offset)
            .expect("a checked probability plan cannot overflow a target offset")
    }

    fn for_each_group(
        &self,
        input: &TensorView<'_>,
        finiteness: LogitFiniteness,
        mut visit: impl FnMut(usize, usize, RowStats) -> Result<(), ProbabilityError>,
    ) -> Result<(), ProbabilityError> {
        for (group, group_base) in self.group_offsets(input).enumerate() {
            let stats = row_stats(input, self, finiteness, group, group_base)?;
            visit(group, group_base, stats)?;
        }
        Ok(())
    }
}

fn checked_finite_logit(value: f64, group: usize, class: usize) -> Result<f64, ProbabilityError> {
    if value.is_nan() {
        Err(ProbabilityError::NaNLogit { group, class })
    } else if value == f64::INFINITY {
        Err(ProbabilityError::PositiveInfinityLogit { group, class })
    } else if value == f64::NEG_INFINITY {
        Err(ProbabilityError::NegativeInfinityLogit { group, class })
    } else {
        Ok(value)
    }
}

fn validate_finite_logits(
    input: &TensorView<'_>,
    plan: &AxisPlan,
) -> Result<FiniteLogits, ProbabilityError> {
    for (group, group_base) in plan.group_offsets(input).enumerate() {
        let mut input_offset = group_base;
        for class in 0..plan.classes {
            checked_finite_logit(input.value_at_storage_offset(input_offset), group, class)?;
            if class + 1 < plan.classes {
                input_offset = input_offset
                    .checked_add(plan.class_stride)
                    .expect("a checked probability plan cannot overflow along the class axis");
            }
        }
    }
    Ok(FiniteLogits)
}

fn row_stats(
    input: &TensorView<'_>,
    plan: &AxisPlan,
    finiteness: LogitFiniteness,
    group: usize,
    group_base: usize,
) -> Result<RowStats, ProbabilityError> {
    debug_assert!(plan.classes > 0);
    let mut maximum = f64::NEG_INFINITY;
    let mut input_offset = group_base;
    for class in 0..plan.classes {
        let value = input.value_at_storage_offset(input_offset);
        let value = match finiteness {
            LogitFiniteness::Check => checked_finite_logit(value, group, class)?,
            LogitFiniteness::Validated(_) => value,
        };
        maximum = maximum.max(value);
        if class + 1 < plan.classes {
            input_offset = input_offset
                .checked_add(plan.class_stride)
                .expect("a checked probability plan cannot overflow along the class axis");
        }
    }

    let mut exponential_tail = 0.0;
    let mut skipped_one_maximum = false;
    input_offset = group_base;
    for class in 0..plan.classes {
        let value = input.value_at_storage_offset(input_offset);
        let shifted = value - maximum;
        if shifted == 0.0 && !skipped_one_maximum {
            skipped_one_maximum = true;
        } else {
            exponential_tail += shifted.exp();
        }
        if class + 1 < plan.classes {
            input_offset = input_offset
                .checked_add(plan.class_stride)
                .expect("a checked probability plan cannot overflow along the class axis");
        }
    }
    debug_assert!(skipped_one_maximum);

    Ok(RowStats {
        maximum,
        shifted_exponential_sum: 1.0 + exponential_tail,
        log_shifted_exponential_sum: exponential_tail.ln_1p(),
    })
}

One forward request creates one checked axis-and-group plan and invokes the row-statistics calculation exactly once for each group. Its emitter then uses those facts to produce log-sum-exp, softmax, log-softmax, or indexed-NLL output. The crate-private log_softmax_forward and indexed_mean_nll_forward helpers may emit probabilities alongside their primary result so a later backward gradient calculation can reuse them without normalizing the logits again. The public functions do not expose the optional saved tensor; each returns only its documented result.

When a helper requests the optional saved tensor, it first checks every logit for a non-finite value before reserving that tensor’s storage. This preserves the same error order as computing the public result before a separate probability tensor. The one row-statistics calculation then performs its maximum and shifted-exponential scans. The preliminary finite-input scan does not calculate either group statistic. A successful preliminary scan returns a private FiniteLogits marker. Because the tensor view cannot mutate its values, the maximum scan trusts that marker instead of checking the same values again.

“Once for each group” does not mean that each logit is read only once. Stable row statistics still require a maximum scan and a shifted-exponential scan, and producing class-wise output requires another class scan. It also does not combine separate public calls: calling softmax and then log_softmax remains two independent forward requests.

A slice or transposed view supplies its own group strides, class stride, and base offset, so the same traversal reads its logical groups without copying them. Softmax and log-softmax create a separate checked group-base cursor for their contiguous output. The input and output cursors visit groups in the same row-major order; within a group, separate input and output class strides place every result at its logical row-major output position. No numerical pass constructs a class coordinate vector or calls TensorView::get once per scalar. Scalar reads and writes still use ordinary safe bounds-checked indexing.

Softmax divides shifted exponentials by $S$ . Log-softmax computes the shifted logit minus $\ln S$ . For $T$ targets, fused indexed mean NLL keeps two accumulators. The ordinary total adds each complete nonnegative group loss. If every group loss and the running sum remain finite, the function divides total by $T$ once and returns the mean in nats per target; this single final division preserves representable subnormal mean rounding.

In parallel, the fallback scaled_mean adds the two nonnegative parts of each group loss after dividing each part by $T$ : the target-logit gap $(m-\ell_{t_r})/T$ , where $t_r$ is the target class for group $r$ , and the log-normalizer $\ln(1+\mathrm{tail})/T$ . If $m-\ell_{t_r}$ itself overflows, the fallback computes $m/T-\ell_{t_r}/T$ instead. The function returns scaled_mean only when a complete group loss or the running value of total overflows; otherwise it divides total by $T$ and returns that quotient. After all target bounds pass, the selected target logit is read at the group’s base plus the target index times the input class stride. That keeps target validation ahead of every logit read and retains group-major error order.

Emit requested probability results from one checked group traversal rust/crates/llm-from-scratch/src/nn/probability.rs#stable-probability-operations

fn output_buffer(elements: usize) -> Result<Vec<f64>, ProbabilityError> {
    let mut values = Vec::new();
    values
        .try_reserve_exact(elements)
        .map_err(|_| ProbabilityError::OutputAllocationFailed { elements })?;
    values.resize(elements, 0.0);
    Ok(values)
}

fn positive_zero(value: f64) -> f64 {
    if value == 0.0 { 0.0 } else { value }
}

/// Requested normalized values emitted by one checked forward traversal.
#[derive(Debug)]
struct NormalizedForward {
    probabilities: Option<Tensor>,
    log_probabilities: Option<Tensor>,
}

/// Log-softmax output and optional probabilities from the same forward traversal.
#[derive(Debug)]
pub(crate) struct LogSoftmaxForward {
    pub(crate) value: Tensor,
    pub(crate) probabilities: Option<Tensor>,
}

/// Indexed mean NLL and optional probabilities emitted by its forward traversal.
#[derive(Debug)]
pub(crate) struct IndexedMeanNllForward {
    pub(crate) loss: f64,
    pub(crate) probabilities: Option<Tensor>,
}

struct NormalizedGroupOutput<'a> {
    output_group_base: usize,
    output_class_stride: usize,
    probabilities: Option<&'a mut [f64]>,
    log_probabilities: Option<&'a mut [f64]>,
}

impl NormalizedGroupOutput<'_> {
    fn emit(
        mut self,
        input: &TensorView<'_>,
        plan: &AxisPlan,
        input_group_base: usize,
        stats: RowStats,
    ) {
        debug_assert!(self.probabilities.is_some() || self.log_probabilities.is_some());
        let mut input_offset = input_group_base;
        let mut output_offset = self.output_group_base;
        for class in 0..plan.classes {
            let shifted = input.value_at_storage_offset(input_offset) - stats.maximum;
            if let Some(values) = self.probabilities.as_mut() {
                values[output_offset] =
                    positive_zero(shifted.exp() / stats.shifted_exponential_sum);
            }
            if let Some(values) = self.log_probabilities.as_mut() {
                values[output_offset] = positive_zero(shifted - stats.log_shifted_exponential_sum);
            }

            if class + 1 < plan.classes {
                input_offset = input_offset
                    .checked_add(plan.class_stride)
                    .expect("a checked probability plan cannot overflow along the class axis");
                output_offset = output_offset
                    .checked_add(self.output_class_stride)
                    .expect("a checked probability output cannot overflow along the class axis");
            }
        }
    }
}

/// Reduces one axis with max-shifted log-sum-exp.
///
/// An empty selected axis returns the log-additive identity, negative infinity,
/// once per remaining-axis group. Other non-finite logits are rejected in
/// group-major, class-minor order.
pub fn log_sum_exp(
    input: &TensorView<'_>,
    axis: usize,
    keep_dim: bool,
) -> Result<Tensor, ProbabilityError> {
    let plan = AxisPlan::new(input, axis, true)?;
    let output_shape = if keep_dim {
        let mut shape = input.shape().to_vec();
        shape[axis] = 1;
        shape
    } else {
        plan.group_shape.clone()
    };
    let (_, output_len) = checked_row_major_layout(&output_shape)?;
    debug_assert_eq!(output_len, plan.groups);
    let mut values = output_buffer(output_len)?;

    if plan.classes == 0 {
        values.fill(f64::NEG_INFINITY);
    } else {
        plan.for_each_group(
            input,
            LogitFiniteness::Check,
            |group, _group_base, stats| {
                values[group] = stats.maximum + stats.log_shifted_exponential_sum;
                Ok(())
            },
        )?;
    }

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

/// Converts finite logits to normalized probabilities along one explicit axis.
pub fn softmax(input: &TensorView<'_>, axis: usize) -> Result<Tensor, ProbabilityError> {
    let forward = normalized_forward(input, axis, true, false)?;
    Ok(forward
        .probabilities
        .expect("softmax requests a probability output"))
}

/// Converts finite logits to normalized log-probabilities along one explicit axis.
pub fn log_softmax(input: &TensorView<'_>, axis: usize) -> Result<Tensor, ProbabilityError> {
    let forward = log_softmax_forward(input, axis, false)?;
    debug_assert!(forward.probabilities.is_none());
    Ok(forward.value)
}

pub(crate) fn log_softmax_forward(
    input: &TensorView<'_>,
    axis: usize,
    emit_probabilities: bool,
) -> Result<LogSoftmaxForward, ProbabilityError> {
    let forward = normalized_forward(input, axis, emit_probabilities, true)?;
    Ok(LogSoftmaxForward {
        value: forward
            .log_probabilities
            .expect("log-softmax forward requests a log-probability output"),
        probabilities: forward.probabilities,
    })
}

fn normalized_forward(
    input: &TensorView<'_>,
    axis: usize,
    emit_probabilities: bool,
    emit_log_probabilities: bool,
) -> Result<NormalizedForward, ProbabilityError> {
    debug_assert!(emit_probabilities || emit_log_probabilities);
    let plan = AxisPlan::new(input, axis, false)?;
    let (output_strides, output_len) = checked_row_major_layout(input.shape())?;
    let mut log_probability_values = emit_log_probabilities
        .then(|| output_buffer(output_len))
        .transpose()?;
    let finiteness = if emit_probabilities && emit_log_probabilities {
        LogitFiniteness::Validated(validate_finite_logits(input, &plan)?)
    } else {
        LogitFiniteness::Check
    };
    let mut probability_values = emit_probabilities
        .then(|| output_buffer(output_len))
        .transpose()?;
    let output_class_stride = output_strides[axis];

    let mut output_group_offsets = plan.output_group_offsets(&output_strides, output_len);
    plan.for_each_group(input, finiteness, |_group, input_group_base, stats| {
        let output_group_base = output_group_offsets
            .next()
            .expect("a checked probability output has one base per input group");
        NormalizedGroupOutput {
            output_group_base,
            output_class_stride,
            probabilities: probability_values.as_deref_mut(),
            log_probabilities: log_probability_values.as_deref_mut(),
        }
        .emit(input, &plan, input_group_base, stats);
        Ok(())
    })?;
    debug_assert!(output_group_offsets.next().is_none());

    let log_probabilities = log_probability_values
        .map(|values| Tensor::from_vec(input.shape().to_vec(), values))
        .transpose()?;
    let probabilities = probability_values
        .map(|values| Tensor::from_vec(input.shape().to_vec(), values))
        .transpose()?;
    Ok(NormalizedForward {
        probabilities,
        log_probabilities,
    })
}

/// Scores one class index per remaining-axis group with fused stable mean NLL.
///
/// Targets follow the row-major group shape obtained by removing `axis` from
/// the logits. Bounds are checked for every target before a logit is read.
pub fn indexed_mean_nll(
    logits: &TensorView<'_>,
    axis: usize,
    targets: &[usize],
) -> Result<f64, ProbabilityError> {
    let forward = indexed_mean_nll_forward(logits, axis, targets, false)?;
    debug_assert!(forward.probabilities.is_none());
    Ok(forward.loss)
}

pub(crate) fn indexed_mean_nll_forward(
    logits: &TensorView<'_>,
    axis: usize,
    targets: &[usize],
    emit_probabilities: bool,
) -> Result<IndexedMeanNllForward, ProbabilityError> {
    let plan = AxisPlan::new(logits, axis, false)?;
    if targets.len() != plan.groups {
        return Err(ProbabilityError::TargetCountMismatch {
            expected: plan.groups,
            actual: targets.len(),
        });
    }
    if targets.is_empty() {
        return Err(ProbabilityError::EmptyTargets);
    }
    for (group, &target) in targets.iter().enumerate() {
        if target >= plan.classes {
            return Err(ProbabilityError::TargetOutOfBounds {
                group,
                target,
                classes: plan.classes,
            });
        }
    }

    let finiteness = if emit_probabilities {
        LogitFiniteness::Validated(validate_finite_logits(logits, &plan)?)
    } else {
        LogitFiniteness::Check
    };

    let output_layout = emit_probabilities
        .then(|| checked_row_major_layout(logits.shape()))
        .transpose()?;
    let mut probability_values = output_layout
        .as_ref()
        .map(|(_, output_len)| output_buffer(*output_len))
        .transpose()?;
    let mut output_group_offsets = output_layout
        .as_ref()
        .map(|(output_strides, output_len)| plan.output_group_offsets(output_strides, *output_len));
    let output_class_stride = output_layout
        .as_ref()
        .map(|(output_strides, _)| output_strides[axis]);

    let mut total = 0.0;
    let mut scaled_mean = 0.0;
    let mut needs_scaled_fallback = false;
    let target_count = targets.len() as f64;
    plan.for_each_group(logits, finiteness, |group, group_base, stats| {
        let target = targets[group];
        let target_logit = logits.value_at_storage_offset(plan.target_offset(group_base, target));
        let gap = stats.maximum - target_logit;
        let scaled_gap = if gap.is_finite() {
            gap / target_count
        } else {
            stats.maximum / target_count - target_logit / target_count
        };
        scaled_mean += scaled_gap + stats.log_shifted_exponential_sum / target_count;

        let loss = gap + stats.log_shifted_exponential_sum;
        if loss.is_finite() && !needs_scaled_fallback {
            total += loss;
            if !total.is_finite() {
                needs_scaled_fallback = true;
            }
        } else {
            needs_scaled_fallback = true;
        }

        if let Some(values) = probability_values.as_deref_mut() {
            let output_group_base = output_group_offsets
                .as_mut()
                .and_then(Iterator::next)
                .expect("a checked probability output has one base per input group");
            NormalizedGroupOutput {
                output_group_base,
                output_class_stride: output_class_stride
                    .expect("a requested probability output has a class stride"),
                probabilities: Some(values),
                log_probabilities: None,
            }
            .emit(logits, &plan, group_base, stats);
        }
        Ok(())
    })?;
    debug_assert!(
        output_group_offsets
            .as_mut()
            .is_none_or(|offsets| offsets.next().is_none())
    );

    let loss = positive_zero(if needs_scaled_fallback {
        scaled_mean
    } else {
        total / target_count
    });
    let probabilities = probability_values
        .map(|values| Tensor::from_vec(logits.shape().to_vec(), values))
        .transpose()?;
    Ok(IndexedMeanNllForward {
        loss,
        probabilities,
    })
}

The example helper constructs all three constant-shift rows once and calls the cumulative API. It does not keep a second implementation of the stable formula.

Run every checked probability operation over the shared three-row example rust/demos/ch12-stable-softmax/src/lib.rs#tiny-stable-softmax-example

/// Normalizes the same relative logits after three different constant shifts.
pub fn tiny_stable_softmax_example() -> Result<TinyStableSoftmaxExample, ProbabilityError> {
    let logits = Tensor::from_vec(LOGIT_SHAPE.to_vec(), LOGIT_VALUES.to_vec())?;
    let probabilities = softmax(&logits.view(), CLASS_AXIS)?;
    let log_probabilities = log_softmax(&logits.view(), CLASS_AXIS)?;
    let log_normalizers = log_sum_exp(&logits.view(), CLASS_AXIS, false)?;
    let mean_nll = indexed_mean_nll(&logits.view(), CLASS_AXIS, &TARGETS)?;

    Ok(TinyStableSoftmaxExample {
        logits,
        probabilities,
        log_probabilities,
        log_normalizers,
        mean_nll,
    })
}

An empty selected axis has one useful log-domain identity: log-sum-exp returns negative infinity for every remaining group. No normalized distribution exists there, so the other operations reject it. A zero extent on a different axis produces a valid empty tensor without reads. Exact zero is canonicalized to positive zero, making singleton log-softmax and NLL exactly positive zero.

Maximum shifting removes avoidable failures, not the limits of f64. An extremely unlikely finite class may still round to probability zero, while its log-probability and fused target loss remain finite. A mathematical class range beyond f64 may still become signed infinity, while log-sum-exp at the upper boundary may round to f64::MAX. Chapter 7 correctly scores an already-rounded zero assigned probability as infinite NLL; this chapter’s fused logit path avoids that rounding when the log-domain answer is representable.

The learner program prints ordinary and extreme results plus four errors:

Prepare stable outputs, raw failure statuses, target loss, invariance, and typed errors rust/demos/ch12-stable-softmax/src/main.rs#learner-stable-softmax-output

    let example = tiny_stable_softmax_example()?;
    let row_sums = example
        .probabilities
        .as_slice()
        .chunks_exact(2)
        .map(|row| row.iter().sum())
        .collect::<Vec<f64>>();
    let target_losses = TARGETS
        .iter()
        .enumerate()
        .map(|(row, &target)| -example.log_probabilities.as_slice()[row * 2 + target])
        .collect::<Vec<_>>();
    let ordinary_direct = direct_output_softmax(&[0.0, 1.0]);
    let overflow_direct = direct_output_softmax(&[1000.0, 1001.0]);
    let underflow_direct = direct_output_softmax(&[-1001.0, -1000.0]);

    let axis_error = softmax(&example.logits.view(), 2).unwrap_err();
    let empty_logits = Tensor::from_vec(vec![2, 0], vec![])?;
    let empty_error = softmax(&empty_logits.view(), 1).unwrap_err();
    let nonfinite_logits = Tensor::from_vec(vec![1, 2], vec![0.0, f64::INFINITY])?;
    let nonfinite_error = softmax(&nonfinite_logits.view(), 1).unwrap_err();
    let target_error = indexed_mean_nll(&example.logits.view(), 1, &[1, 2, 1]).unwrap_err();

./course run cargo run --quiet --locked -p ch12-stable-softmax

stable softmax: shape=[3, 2] values=[0.268941421370, 0.731058578630, 0.268941421370, 0.731058578630, 0.268941421370, 0.731058578630]
log-sum-exp: shape=[3] values=[1.313261687518, 1001.313261687518, -999.686738312482]
targets: [1, 0, 1] losses=[0.313261687518, 1.313261687518, 0.313261687518] mean_nll=0.646595020852
naive overflow [1000, 1001]: undefined=true
naive underflow [-1001, -1000]: undefined=true

The complete output also reports logits, log-softmax, probability sums, the ordinary direct result, shift equality, exact errors, and the Chapter 13 handoff. Together these records distinguish rounded decimal output from exact shapes, error variants, infinities, and positive-zero behavior.

Compare naive and stable exponentials

The table aligns each direct-path status with the maximum-shifted stages for the same row. It shows the maximum, shifted values, denominator, probabilities, selected target, and loss without recalculating the Rust results. Distinct symbols and border patterns keep finite, stable, overflow, underflow, and rejected states separate without relying on color.

See one maximum shift rescue ordinary and extreme logits

Compare three Rust-recorded rows with equal relative logits, follow their stable normalization and target losses, and inspect four rejected requests.

Logit shape: [3, 2]
Class axis: 1
Mean NLL: 0.646595020852

Subtract one row maximum before exponentiating

Each row begins at a different absolute scale. The table shows that subtracting its maximum produces the same shifted values, exponentials, and normalized probabilities.

Subtract one row maximum before exponentiating	Row 0	Row 1	Row 2
Raw logits	[0.000000000000, 1.000000000000]	[1000.000000000000, 1001.000000000000]	[-1001.000000000000, -1000.000000000000]
Maximum	1.000000000000	1001.000000000000	-1000.000000000000
Shifted logits	[-1.000000000000, 0.000000000000]	[-1.000000000000, 0.000000000000]	[-1.000000000000, 0.000000000000]
Exponentials	[0.367879441171, 1.000000000000]	[0.367879441171, 1.000000000000]	[0.367879441171, 1.000000000000]
Raw-exponential path	finite : [1.000000000000, 2.718281828459]	undefined after overflow	undefined after underflow
Probabilities	[0.268941421370, 0.731058578630]	[0.268941421370, 0.731058578630]	[0.268941421370, 0.731058578630]

All three recorded rows have exactly matching probabilities. Denominator: 1.367879441171

Select one log-probability per target

For each row, the target index selects the recorded log-probability. Its negation is the row loss; the three row losses average to the mean NLL shown above.

Row 0

Class: 1

Log-probabilities: -0.313261687518

Negative log-likelihood: 0.313261687518

Row 1

Class: 0

Log-probabilities: -1.313261687518

Negative log-likelihood: 1.313261687518

Row 2

Class: 1

Log-probabilities: -0.313261687518

Negative log-likelihood: 0.313261687518

Reject invalid axes, logits, and targets

Dashed rejection cards preserve the exact axis, group, class, and target evidence returned before an invalid computation continues.

`axis-out-of-bounds`

Class axis 2; Rank 2

`empty-normalization-axis`

Class axis 1

`positive-infinity-logit`

Group 0; Class 1

`target-out-of-bounds`

Group 1; Class 2; Classes 2

Predict before running Rust

Predict the shifted values for [1000,1001].
Predict whether adding -1001 to [0,1] changes either probability.
Explain why direct normalization of [1000,1001] is undefined in f64.
Predict both probabilities for equal logits [7,7].
Select the loss for target class 0 in row [1000,1001].
Predict log-sum-exp output shapes for input [2,3,4], axis 1, with and without keep_dim.
Decide which empty-class operation has a defined identity.
For shape [3,2], source strides [2,1], and class axis 1, list the three group-base offsets and the two source offsets read in each group.
Suppose one training operation must return log-softmax values and retain softmax probabilities for its backward gradient calculation. Which group-wide facts can both results share, and which class-wise work still remains?
Misconception check: does maximum shifting itself turn logits into probabilities?

Check the predictions

The maximum is 1001, so the shifted values are exactly [-1,0].
No. Adding one shared constant changes the maximum by the same amount and leaves both shifted differences unchanged.
Both raw exponentials overflow to infinity, so each division becomes infinity divided by infinity rather than a probability.
Equal shifted logits have equal exponentials, so both probabilities are 0.5.
Class zero has log-probability -1.313261687518, so its NLL is 1.313261687518.
Removing axis one gives [2,4]; retaining it gives [2,1,4].
Log-sum-exp returns the log-additive identity negative infinity. A softmax distribution, log-softmax distribution, and indexed target loss need at least one class.
Removing class axis 1 leaves group stride [2], so the group bases are [0,2,4]. Class stride 1 gives source offsets [0,1], [2,3], and [4,5].
Both results share the already computed maximum $m$ , shifted-exponential sum $S$ , and log-normalizer $\ln S$ . One class scan can emit both values from those facts. The operation does not repeat the maximum or shifted-sum calculation, but it still visits every class to write the requested outputs.
No. The shift only stabilizes relative logits. Exponentiation and division by the complete shifted sum produce probabilities.

Run the example after predicting:

./course run cargo run --quiet --locked -p ch12-stable-softmax

Prepare a sampled gradient cross-check

The cumulative tensor core can now turn finite strided logits into owned probabilities, log-probabilities, log-sum-exp values, and fused indexed mean NLL along any explicit axis. These operations will normalize vocabulary and attention scores and provide the forward indexed mean NLL that Chapter 13 uses for a materially separate sampled finite-difference cross-check. The analytic and numerical paths still share the fixture logits and target indices, IEEE f64 arithmetic and its elementary exp, Tensor storage, and row-major index conventions, so agreement is evidence for the selected probes of a locally smooth objective, not proof of the complete gradient or every shared assumption.

Inside those operations, one checked cursor identifies each normalization group and one class stride selects the group’s logits. One forward request computes the maximum, shifted-exponential sum, and log-normalizer once per group; a training request can retain probabilities emitted from those same facts for a future backward gradient calculation. This reuse does not change the probabilities or the loss. Chapter 13 perturbs that production indexed_mean_nll objective, while a separate local analytic routine computes stabilized row probabilities and candidate derivatives without calling the production softmax or indexed_mean_nll implementation.

Finite differences vary one scalar logit at a time and require the objective to be smooth across each probe interval. The two paths use the same fixture logits and target indices, IEEE f64 arithmetic and its elementary exp, Tensor storage, and row-major index conventions. Agreement at the selected coordinates therefore supports only that fixture, those probes, step, and tolerance. It does not prove differentiability, the complete gradient, the complete probability implementation, or any shared assumption.