16 · Content revision 7

Reverse the operations that turn token IDs into loss

Implement VJPs for matrix products, repeated embedding lookups, SiLU, log-softmax, and mean token loss, then compare each new local rule with sampled central differences.

Predict one repeated-token path

Start with a three-row embedding table, four token IDs, a two-class projection, and four target classes:

E=\begin{bmatrix}2&2\\1&-1\\-1&1\end{bmatrix},\qquad z=\begin{bmatrix}1&1&1&2\end{bmatrix},\qquad \mathrm{targets}=\begin{bmatrix}0&0&0&1\end{bmatrix}

W=\begin{bmatrix}1&-1\\1&-1\end{bmatrix}

The formula below names occurrences by batch and token coordinates $(b,t)$ . This example stores them in row-major order as flat positions $p=bT+t$ . With $B=1$ and $T=4$ , positions $0,1,2,3$ correspond to $(0,0),(0,1),(0,2),(0,3)$ .

Row gather materializes [1,-1] three times and [-1,1] once. Every selected row multiplied by $W$ yields projection preactivation [0,0]. SiLU keeps zero. For this compact exercise, those activated values are used directly as two-class loss logits, so log-softmax produces $[-\ln 2,-\ln 2]$ . Every target therefore has loss $\ln 2$ , and their mean is also $\ln 2$ .

This chain is deliberately an operation exercise, not a decoder architecture. In the later decoder, SiLU belongs inside a feed-forward block; a separate vocabulary projection produces the logits consumed by token loss.

The displayed log-softmax and the combined indexed mean NLL are two separate operation calls. They share the input logits, not a forward result. Each call saves its own forward evidence for its later VJP.

Predict the reverse signs before calculating values. The correct-target logit must move up to reduce loss, so its loss gradient is negative; the competing logit gradient is positive. With two equal classes and four mean groups, their magnitudes are $1/8$ :

p\in\{0,1,2\},\ \mathrm{target}_p=0: \quad\begin{bmatrix}-\frac18&\frac18\end{bmatrix};\qquad p=3,\ \mathrm{target}_3=1: \quad\begin{bmatrix}\frac18&-\frac18\end{bmatrix}

SiLU’s derivative at zero is $1/2$ . The projection VJP then gives three occurrence gradients $[-1/8,-1/8]$ for ID 1 and one $[1/8,1/8]$ for ID 2. Do not average the three ID 1 contributions again: the $1/4$ loss mean is already inside each one.

dE= \begin{bmatrix} 0 & 0 \\ -\frac{3}{8} & -\frac{3}{8} \\ \frac{1}{8} & \frac{1}{8} \end{bmatrix}, \qquad dW= \begin{bmatrix} -\frac{1}{4} & \frac{1}{4} \\ \frac{1}{4} & -\frac{1}{4} \end{bmatrix}

The fixed Rust calculation makes this compact chain explicit before the same work is expressed through reusable local operations:

Compute one fixed token-to-loss backward path before introducing reusable model-operation VJPs rust/demos/ch16-model-autodiff-ops/src/lib.rs#handwritten-model-backward

/// Computes the compact token-operation chain with fixed arrays and handwritten rules.
///
/// This bounded reference calculation illustrates the model-specific backward
/// style that preceded a reusable operation vocabulary; it is not attributed
/// source code from any historical paper.
pub fn handwritten_model_backward() -> HandwrittenModelBaseline {
    let mut gathered = [0.0; 8];
    for (position, &token_id) in TOKEN_IDS.iter().enumerate() {
        for feature in 0..EMBEDDING_SHAPE[1] {
            gathered[position * 2 + feature] =
                EMBEDDING_VALUES[token_id * EMBEDDING_SHAPE[1] + feature];
        }
    }

    let mut projection_preactivations = [0.0; 8];
    for position in 0..TOKEN_IDS.len() {
        for class in 0..2 {
            for feature in 0..2 {
                projection_preactivations[position * 2 + class] +=
                    gathered[position * 2 + feature] * WEIGHT_VALUES[feature * 2 + class];
            }
        }
    }

    let mut activated = [0.0; 8];
    let mut log_probabilities = [0.0; 8];
    let mut loss = 0.0;
    let mut loss_input_gradient = [0.0; 8];
    let mut matmul_output_gradient = [0.0; 8];
    for position in 0..TOKEN_IDS.len() {
        for class in 0..2 {
            let offset = position * 2 + class;
            activated[offset] = projection_preactivations[offset]
                * stable_sigmoid(projection_preactivations[offset]);
        }
        let maximum = activated[position * 2].max(activated[position * 2 + 1]);
        let denominator = (activated[position * 2] - maximum).exp()
            + (activated[position * 2 + 1] - maximum).exp();
        for class in 0..2 {
            let offset = position * 2 + class;
            let probability = (activated[offset] - maximum).exp() / denominator;
            log_probabilities[offset] = probability.ln();
            let target_indicator = usize::from(TARGETS[position] == class) as f64;
            loss_input_gradient[offset] = (probability - target_indicator) / TOKEN_IDS.len() as f64;
            let sigmoid = stable_sigmoid(projection_preactivations[offset]);
            let silu_derivative =
                sigmoid * (1.0 + projection_preactivations[offset] * (1.0 - sigmoid));
            matmul_output_gradient[offset] = loss_input_gradient[offset] * silu_derivative;
        }
        loss -= log_probabilities[position * 2 + TARGETS[position]] / TOKEN_IDS.len() as f64;
    }

    let mut gathered_gradient = [0.0; 8];
    let mut weight_gradient = [0.0; 4];
    for position in 0..TOKEN_IDS.len() {
        for feature in 0..2 {
            for class in 0..2 {
                gathered_gradient[position * 2 + feature] += matmul_output_gradient
                    [position * 2 + class]
                    * WEIGHT_VALUES[feature * 2 + class];
                weight_gradient[feature * 2 + class] +=
                    gathered[position * 2 + feature] * matmul_output_gradient[position * 2 + class];
            }
        }
    }

    let mut embedding_gradient = [0.0; 6];
    for (position, &token_id) in TOKEN_IDS.iter().enumerate() {
        for feature in 0..2 {
            embedding_gradient[token_id * 2 + feature] += gathered_gradient[position * 2 + feature];
        }
    }

    HandwrittenModelBaseline {
        gathered,
        projection_preactivations,
        activated,
        log_probabilities,
        loss,
        loss_input_gradient,
        matmul_output_gradient,
        gathered_gradient,
        embedding_gradient,
        weight_gradient,
    }
}

Add every occurrence to its shared embedding row

The chapter’s display formula is:

\frac{\partial L}{\partial E_{i,:}}=\sum_{(b,t):z_{b,t}=i}\frac{\partial L}{\partial X_{b,t,:}}

The forward gather result owns four row values. The tape retains the selector for each occurrence, so reverse mode routes each occurrence’s adjoint back to its source row. Repeated selectors therefore accumulate into the same entry of $\bar E$ ; they do not overwrite one another and they do not introduce another mean.

Name the loss, table, selectors, and adjoints

Symbol	Operational meaning
$L$	The scalar mean token loss.
$E$	The trainable embedding table with shape $[V,d]$ .
$i$	One vocabulary-row index in $E$ .
$:$	Every feature coordinate of the named row.
$b$	One batch index.
$t$	One token-position index.
$z_{b,t}$	The integer token ID at that batch-position pair.
$X_{b,t,:}$	The gathered feature row consumed at that occurrence.
$\partial L/\partial X_{b,t,:}$	The upstream occurrence adjoint.
$\sum_{(b,t):z_{b,t}=i}$	Visit every occurrence whose ID equals row $i$ .
$\partial L/\partial E_{i,:}$	The accumulated adjoint for all features of table row $i$ .

Chapter 15 wrote a parent adjoint as $\bar{x}$ . Here the same rule is $\bar E_{i,:}\mathrel{+}=\bar X_{b,t,:}$ for every matching selector. Token IDs do not receive gradients; they choose destinations. If occurrences are stored flat, $p=bT+t$ maps each flat position back to the same logical pair used in the conditioned sum.

From one neural next-word backward pass to reusable decoder VJPs

Bengio et al. train a neural next-word model with a learned word-feature table, matrix transforms, a tanh hidden layer, output probabilities, and explicit model-specific backward/update equations. That presentation makes the full learning path inspectable, but Chapter 15’s structural tensor tape still cannot express the lookup, matrix, activation, normalization, and token-loss derivatives needed to train even this small language-model path.

This earlier learned-row path is described by Bengio et al., A Neural Probabilistic Language Model. Bengio et al. build a neural next-word model from learned word-feature rows, matrix equations, a tanh hidden layer, normalized output probabilities, and an explicit backward/update phase for the model parameters.

The fixed calculation above mirrors that inspectable, model-specific style. It is a new Rust example, not code copied from or attributed to the paper. Adding a different operation would require editing its whole backward path.

Abadi et al. describe tensor operation graphs whose differentiation finds every path from a loss to parameters and sums partial-gradient contributions, including gathered embedding rows. Vaswani et al. place learned embeddings and the output projection at model boundaries while matrix projections, softmax attention, and nonlinear feed-forward sublayers repeat through the Transformer stack. Shazeer later evaluates Swish with beta one—the same function as SiLU—and SwiGLU variants inside Transformer feed-forward sublayers.

The transition to reusable operation graphs is described by Abadi et al., TensorFlow: A System for Large-Scale Machine Learning. Abadi et al. represent operations as graph vertices and tensors as edge values, describe automatic differentiation that sums every backward path to a parameter, and show Gather-based embedding graphs whose gradients update gathered rows.

The Transformer arrangement is described by Vaswani et al., Attention Is All You Need. Vaswani et al. construct the Transformer from learned embeddings, learned query/key/value projections, attention softmax, two-transform ReLU feed-forward sublayers, and a learned output transform followed by softmax.

The later gated feed-forward step is described by Shazeer, GLU Variants Improve Transformer. Shazeer defines Swish as its input multiplied by the sigmoid of beta times that input, so beta one gives the function also called SiLU; the paper uses it in SwiGLU Transformer feed-forward variants and reports improved held-out log-perplexity for gated variants over the studied baseline. The original Transformer citation uses ReLU; it is not evidence that Vaswani et al. used SiLU.

This chapter supplies reusable local VJPs for batched matrix products, repeated row gathers, exp, log, SiLU, stable log-softmax, and combined indexed mean NLL. Log-softmax and combined indexed mean NLL each retain the probability values emitted by their own forward call for their VJP instead of normalizing the logits again. These operations form the local reverse rules later embedding, projection, SwiGLU, attention, and token-loss components need; ordinary inference uses only the forward paths.

Save the forward data each local VJP needs

The public methods extend the Chapter 15 tape instead of constructing a second graph. TensorValue::gather_rows is the checked public entry for row selection. The tape first establishes that its table operand is available. For an available operand, the method validates these facts in order:

the table has rank two;
index_shape is a valid row-major logical shape and therefore has a known number of positions;
that position count equals indices.len(); and
each selector, visited in flat order, is smaller than the table’s row count.

The flat-order scan reports the first invalid selector. Only after all four facts hold does the operation append the table width to index_shape and validate the output shape. In the worked example, table shape [3,2], index_shape=[4], and selectors [1,1,1,2] produce output shape [4,2].

One crate-private RowGatherPlan then owns the selectors, logical shape, parent shape, derived output shape, and output element count. Private fields prevent a public caller from constructing a plan that bypasses those checks. The plan establishes dimensions and bounds; allocating the output buffer can still fail separately.

Validate a raw row-gather request once and seal its selectors and shapes in an owned plan rust/crates/llm-from-scratch/src/autograd/model_ops.rs#model-row-gather-plan

/// Owned row-gather facts established before materialization begins.
#[derive(Debug)]
pub(crate) struct RowGatherPlan {
    indices: Vec<usize>,
    index_shape: Vec<usize>,
    input_shape: [usize; 2],
    output_shape: Vec<usize>,
    output_len: usize,
}

impl RowGatherPlan {
    fn checked(
        table: &Tensor,
        indices: &[usize],
        index_shape: &[usize],
    ) -> Result<Self, TensorAutodiffError> {
        if table.rank() != 2 {
            return Err(ModelOpError::GatherTableRank { rank: table.rank() }.into());
        }
        let (_, expected) = checked_row_major_layout(index_shape)?;
        if indices.len() != expected {
            return Err(ModelOpError::GatherIndexCountMismatch {
                expected,
                actual: indices.len(),
            }
            .into());
        }
        let rows = table.shape()[0];
        for (position, &index) in indices.iter().enumerate() {
            if index >= rows {
                return Err(ModelOpError::GatherIndexOutOfBounds {
                    position,
                    index,
                    rows,
                }
                .into());
            }
        }

        Self::from_validated_indices(table, indices.to_vec(), index_shape.to_vec())
    }

    /// Seals indices whose shape, count, and bounds were established by a
    /// crate-owned caller for this exact rank-two table.
    pub(crate) fn from_validated_indices(
        table: &Tensor,
        indices: Vec<usize>,
        index_shape: Vec<usize>,
    ) -> Result<Self, TensorAutodiffError> {
        let input_shape = [table.shape()[0], table.shape()[1]];
        let width = input_shape[1];
        let mut output_shape = index_shape.clone();
        output_shape
            .try_reserve_exact(1)
            .map_err(|_| ModelOpError::OutputAllocationFailed {
                elements: indices.len().saturating_mul(width),
            })?;
        output_shape.push(width);
        let (_, output_len) = checked_row_major_layout(&output_shape)?;
        Ok(Self {
            indices,
            index_shape,
            input_shape,
            output_shape,
            output_len,
        })
    }

    fn into_saved_context(self) -> ModelSavedContext {
        let Self {
            indices,
            index_shape,
            input_shape,
            output_shape,
            ..
        } = self;
        ModelSavedContext::GatherRows {
            indices,
            index_shape,
            input_shape: input_shape.to_vec(),
            output_shape,
        }
    }
}

The forward row-copy kernel receives the plan rather than raw selectors and shapes, so it does not rescan rank, count, or bounds. It materializes independent row values; the result does not alias the table. After copying succeeds, the plan’s selectors and shapes move into the saved gather context. “Trusted” here means that the kernel consumes a value whose private construction established the facts. It does not mean that public input may be unchecked.

Create the checked plan after operand validation, copy its rows, and retain its facts for reversal rust/crates/llm-from-scratch/src/autograd/model_ops.rs#model-row-gather-operation

    /// Selects rows from a rank-two table into `index_shape + [width]`.
    ///
    /// IDs are integer selectors and are deliberately not tape operands.
    pub fn gather_rows(
        &self,
        indices: &[usize],
        index_shape: &[usize],
    ) -> Result<Self, TensorAutodiffError> {
        self.gather_rows_with_plan(|table| RowGatherPlan::checked(table, indices, index_shape))
    }

    /// Builds one row-gather plan after operand availability is established.
    pub(crate) fn gather_rows_with_plan(
        &self,
        build_plan: impl FnOnce(&Tensor) -> Result<RowGatherPlan, TensorAutodiffError>,
    ) -> Result<Self, TensorAutodiffError> {
        Self::model_operation(TensorOperation::GatherRows, [self], |primals| {
            let table = primals[0];
            let plan = build_plan(table)?;
            let value = gather_rows_forward(table, &plan)?;
            Ok((value, [plan.into_saved_context()]))
        })
    }

Gather reversal allocates one zero tensor with the parent table’s shape. For each flat occurrence, it adds the upstream feature row to the destination named by the saved selector. This is where repeated IDs meet: the forward output rows are separate values, but their reverse contributions can share a destination.

Add each occurrence adjoint into the embedding-table row selected in the forward pass rust/crates/llm-from-scratch/src/autograd/model_ops.rs#model-row-gather-vjp

        ModelSavedContext::GatherRows {
            indices,
            input_shape,
            output_shape,
            ..
        } => {
            debug_assert_eq!(upstream.shape(), output_shape);
            let width = input_shape[1];
            let mut table_gradient = zeros(input_shape)?;
            for (position, &index) in indices.iter().enumerate() {
                let source = position * width;
                let destination = index * width;
                for feature in 0..width {
                    table_gradient.as_mut_slice()[destination + feature] +=
                        upstream.as_slice()[source + feature];
                }
            }
            Ok(table_gradient)
        }

The other rules follow the same local pattern. Matrix VJPs transpose the final matrix axes and sum expanded batch contributions back to each parent shape. exp multiplies by its saved output; log divides by its saved positive input; SiLU uses its saved input and stable sigmoid.

TensorValue::log_softmax makes one checked forward call for two outputs: the log-probability tensor returned to the caller and the ordinary probability tensor needed by its VJP. After input validation, one invocation of the checked group driver computes one statistics bundle per group and emits both tensors from those shared statistics. The operation saves the emitted probabilities; backward does not call softmax or normalize the logits again.

Return log-probabilities and retain probabilities from one checked forward call rust/crates/llm-from-scratch/src/autograd/model_ops.rs#model-log-softmax-saved-forward

    /// Applies stable log-softmax along one explicit class axis.
    pub fn log_softmax(&self, axis: usize) -> Result<Self, TensorAutodiffError> {
        Self::model_operation(TensorOperation::LogSoftmax, [self], |primals| {
            let input = primals[0];
            let forward = log_softmax_forward(&input.view(), axis, true)?;
            let probabilities = forward
                .probabilities
                .expect("the autodiff log-softmax forward requests saved probabilities");
            Ok((
                forward.value,
                [ModelSavedContext::LogSoftmax {
                    probabilities,
                    axis,
                    input_shape: input.shape().to_vec(),
                }],
            ))
        })
    }

TensorValue::indexed_mean_nll independently makes one checked forward call for the scalar mean loss and the probabilities needed by its VJP. After the tape establishes operand availability, the operation validates the axis, class extent, target count, nonempty target set, and every target. The probability module checks every logit for finiteness before allocating the optional probability tensor so the established error order is preserved. “One forward call” does not mean one read of each logit: the validation scan computes no maximum, exponential, sum, or output. The checked group driver then calculates one RowStats bundle per group with one class scan for the maximum and a separate class scan for the shifted-exponential sum. The group callback uses that bundle to accumulate the scalar loss, while the emitter scans the classes to write the probability tensor retained for the VJP. The saved tensor contains the same emitted f64 values, bit for bit. This is identity of the stored floating-point values, not a claim of exact real-number arithmetic. Those scans are the steps of one normalization calculation, not a second normalization call.

Return mean NLL and retain probabilities from one checked forward call rust/crates/llm-from-scratch/src/autograd/model_ops.rs#model-indexed-nll-saved-forward

    /// Computes one stable rank-zero mean NLL from flat group-major targets.
    pub fn indexed_mean_nll(
        &self,
        axis: usize,
        targets: &[usize],
    ) -> Result<Self, TensorAutodiffError> {
        Self::model_operation(TensorOperation::IndexedMeanNll, [self], |primals| {
            let logits = primals[0];
            let forward = indexed_mean_nll_forward(&logits.view(), axis, targets, true)?;
            let probabilities = forward
                .probabilities
                .expect("the autodiff indexed-NLL forward requests saved probabilities");
            let value = Tensor::from_vec(Vec::new(), vec![forward.loss])?;
            Ok((
                value,
                [ModelSavedContext::IndexedMeanNll {
                    probabilities,
                    targets: targets.to_vec(),
                    axis,
                    input_shape: logits.shape().to_vec(),
                    groups: targets.len(),
                }],
            ))
        })
    }

The explicit log-softmax and indexed-NLL branches in this lesson remain two operation calls. Each owns its own returned value and saved context. The reuse is between one operation’s forward calculation and its later VJP, not between the two calls.

Log-softmax reversal subtracts probability times the upstream class-axis sum. For group $g$ , class $c$ , input logit $Z_{g,c}$ , target $y_g$ , $G$ groups, saved probability $P_{g,c}$ , and incoming scalar adjoint $\bar L$ , the combined indexed mean NLL gives the logit adjoint

\bar Z_{g,c}=\frac{\bar L}{G}\left(P_{g,c}-\mathbf{1}[c=y_g]\right).

At the target class, the parenthesized term subtracts one from the saved probability. Every class component is then multiplied by the same upstream and mean scale $\bar L/G$ .

The complete example builds gather, matmul, SiLU, displayed log-softmax, and combined loss from the values above, then compares its forward and reverse tensors with the fixed-array calculation. It calls backward_with_trace because that comparison reads the adjoints of the SiLU output, matrix-product output, and gathered-row output. Ordinary training uses a lean method: backward when implicit graph retention is appropriate, or backward_with_seed when the caller must choose retention or release. The lean and traced methods use the same reverse calculation; tracing records its results rather than implementing the local VJPs a second time.

Build the repeated-token operation chain on the shared tape and compare it with the fixed calculation rust/demos/ch16-model-autodiff-ops/src/lib.rs#shared-model-vjp-fixture

/// Runs gather, matmul, SiLU, stable log-softmax, and combined token loss on one
/// repeated-token example, then checks the tape against the fixed reference.
pub fn frozen_model_example() -> Result<FrozenModelExample, TensorAutodiffError> {
    let embeddings = TensorValue::parameter(tensor(&EMBEDDING_SHAPE, &EMBEDDING_VALUES))?;
    let gathered = embeddings.gather_rows(&TOKEN_IDS, &TOKEN_SHAPE)?;
    let weights = TensorValue::parameter(tensor(&WEIGHT_SHAPE, &WEIGHT_VALUES))?;
    let projection_preactivations = gathered.matmul(&weights)?;
    let activated = projection_preactivations.silu()?;
    let log_probabilities = activated.log_softmax(CLASS_AXIS)?;
    let loss = activated.indexed_mean_nll(CLASS_AXIS, &TARGETS)?;
    let baseline = handwritten_model_backward();

    assert_close(gathered.value().as_slice(), &baseline.gathered, 1e-12);
    assert_close(
        projection_preactivations.value().as_slice(),
        &baseline.projection_preactivations,
        1e-12,
    );
    assert_close(activated.value().as_slice(), &baseline.activated, 1e-12);
    assert_close(
        log_probabilities.value().as_slice(),
        &baseline.log_probabilities,
        1e-12,
    );
    assert_close(loss.value().as_slice(), &[baseline.loss], 1e-12);

    let backward = loss.backward_with_trace()?;
    let loss_input_gradient = pass_adjoint(&backward, TensorOperation::Silu);
    let matmul_output_gradient = pass_adjoint(&backward, TensorOperation::MatMul);
    let gathered_gradient = pass_adjoint(&backward, TensorOperation::GatherRows);
    let embedding_gradient = embeddings
        .gradient_snapshot()
        .expect("embeddings are a parameter");
    let weight_gradient = weights
        .gradient_snapshot()
        .expect("weights are a parameter");
    assert_close(
        loss_input_gradient.as_slice(),
        &baseline.loss_input_gradient,
        1e-12,
    );
    assert_close(
        matmul_output_gradient.as_slice(),
        &baseline.matmul_output_gradient,
        1e-12,
    );
    assert_close(
        gathered_gradient.as_slice(),
        &baseline.gathered_gradient,
        1e-12,
    );
    assert_close(
        embedding_gradient.as_slice(),
        &baseline.embedding_gradient,
        1e-12,
    );
    assert_close(weight_gradient.as_slice(), &baseline.weight_gradient, 1e-12);

    Ok(FrozenModelExample {
        embeddings: embeddings.value_snapshot(),
        token_ids: TOKEN_IDS.to_vec(),
        weights: weights.value_snapshot(),
        targets: TARGETS.to_vec(),
        gathered: gathered.value_snapshot(),
        projection_preactivations: projection_preactivations.value_snapshot(),
        activated: activated.value_snapshot(),
        log_probabilities: log_probabilities.value_snapshot(),
        loss: loss.value_snapshot(),
        backward,
        loss_input_gradient,
        matmul_output_gradient,
        gathered_gradient,
        embedding_gradient,
        weight_gradient,
        baseline,
    })
}

Each new VJP is compared with sampled central differences. The two matrix parents are checked separately; gather, exp, log, SiLU, log-softmax, and indexed mean NLL each get their own selected coordinates.

Compare every new local VJP with sampled central differences rust/demos/ch16-model-autodiff-ops/src/lib.rs#model-vjp-gradchecks

/// Checks both matmul parents plus gather, exp, log, SiLU, log-softmax, and NLL.
pub fn model_vjp_gradchecks() -> Result<ModelVjpGradchecks, Box<dyn Error>> {
    let left = tensor(&[2, 2], &[0.4, -0.7, 1.2, 0.3]);
    let right = tensor(&[2, 2], &[0.2, -0.4, 0.9, 0.5]);
    let right_for_tape = right.clone();
    let right_for_objective = right.clone();
    let matmul_left = sampled_model_check(
        left.clone(),
        move |parameter| {
            let right = TensorValue::constant(right_for_tape)?;
            sum_to_scalar(parameter.matmul(&right)?)
        },
        move |candidate| {
            tensor_matmul(&candidate.view(), &right_for_objective.view())
                .expect("the worked matmul shapes agree")
                .as_slice()
                .iter()
                .sum()
        },
    )?;
    let left_for_tape = left.clone();
    let left_for_objective = left.clone();
    let matmul_right = sampled_model_check(
        right,
        move |parameter| {
            let left = TensorValue::constant(left_for_tape)?;
            sum_to_scalar(left.matmul(parameter)?)
        },
        move |candidate| {
            tensor_matmul(&left_for_objective.view(), &candidate.view())
                .expect("the worked matmul shapes agree")
                .as_slice()
                .iter()
                .sum()
        },
    )?;

    let gather_ids = [2, 1, 2];
    let gather = sampled_model_check(
        tensor(&[3, 2], &[0.2, -0.4, 0.7, 1.1, -0.3, 0.6]),
        |parameter| sum_to_scalar(parameter.gather_rows(&gather_ids, &[3])?),
        |candidate| raw_gather_sum(candidate, &gather_ids),
    )?;

    let exp = sampled_model_check(
        tensor(&[3], &[-0.8, 0.2, 1.1]),
        |parameter| sum_to_scalar(parameter.exp()?),
        |candidate| candidate.as_slice().iter().map(|value| value.exp()).sum(),
    )?;
    let log = sampled_model_check(
        tensor(&[3], &[0.4, 1.1, 2.3]),
        |parameter| sum_to_scalar(parameter.log()?),
        |candidate| candidate.as_slice().iter().map(|value| value.ln()).sum(),
    )?;
    let silu = sampled_model_check(
        tensor(&[3], &[-0.8, 0.2, 1.1]),
        |parameter| sum_to_scalar(parameter.silu()?),
        |candidate| {
            candidate
                .as_slice()
                .iter()
                .map(|&value| value * stable_sigmoid(value))
                .sum()
        },
    )?;

    let log_softmax_weights = tensor(&[2, 3], &[0.2, -0.5, 0.7, 1.1, -0.4, 0.3]);
    let weights_for_tape = log_softmax_weights.clone();
    let weights_for_objective = log_softmax_weights.clone();
    let log_softmax = sampled_model_check(
        tensor(&[2, 3], &[0.7, -0.4, 1.1, -0.2, 0.3, 0.8]),
        move |parameter| {
            let weights = TensorValue::constant(weights_for_tape)?;
            sum_to_scalar(parameter.log_softmax(1)?.mul(&weights)?)
        },
        move |candidate| {
            let output = tensor_log_softmax(&candidate.view(), 1)
                .expect("the worked probability axis is valid");
            weighted_sum(output.as_slice(), weights_for_objective.as_slice())
        },
    )?;
    let indexed_mean_nll = sampled_model_check(
        tensor(&[2, 3], &[0.7, -0.4, 1.1, -0.2, 0.3, 0.8]),
        |parameter| parameter.indexed_mean_nll(1, &[2, 0]),
        |candidate| {
            tensor_indexed_mean_nll(&candidate.view(), 1, &[2, 0])
                .expect("the worked targets are valid")
        },
    )?;

    let checks = vec![
        NamedModelGradcheck {
            operation: "matmul-left",
            report: matmul_left,
        },
        NamedModelGradcheck {
            operation: "matmul-right",
            report: matmul_right,
        },
        NamedModelGradcheck {
            operation: "gather_rows",
            report: gather,
        },
        NamedModelGradcheck {
            operation: "exp",
            report: exp,
        },
        NamedModelGradcheck {
            operation: "log",
            report: log,
        },
        NamedModelGradcheck {
            operation: "silu",
            report: silu,
        },
        NamedModelGradcheck {
            operation: "log_softmax",
            report: log_softmax,
        },
        NamedModelGradcheck {
            operation: "indexed_mean_nll",
            report: indexed_mean_nll,
        },
    ];
    let passed = checks.iter().all(|check| check.report.passed);
    Ok(ModelVjpGradchecks { checks, passed })
}

Invalid token IDs, invalid target classes, empty target sets, and exponential overflow are rejected without changing previously recorded gradients. The example snapshots the affected gradient immediately before and after each failure rather than relying only on an initially zero state:

Reject invalid selectors, targets, empty means, and overflow while preserving recorded gradients rust/demos/ch16-model-autodiff-ops/src/lib.rs#model-op-errors-example

pub fn model_error_example() -> Result<ModelErrorExample, TensorAutodiffError> {
    let table = TensorValue::parameter(tensor(&[2, 2], &[1.0, 2.0, 3.0, 4.0]))?;
    let logits = TensorValue::parameter(tensor(&[2, 2], &[1.0, -1.0, 0.5, -0.5]))?;
    let empty_logits = TensorValue::parameter(tensor(&[0, 2], &[]))?;
    let overflow = TensorValue::parameter(tensor(&[], &[f64::MAX]))?;

    sum_to_scalar(table.clone())?.backward()?;
    sum_to_scalar(logits.clone())?.backward()?;
    overflow.backward()?;

    let invalid_id_before = tensor_bits(&table.gradient().expect("table gradient exists"));
    let invalid_id = table
        .gather_rows(&[2], &[1])
        .expect_err("row two is outside a two-row table");
    let invalid_id_gradient_unchanged =
        invalid_id_before == tensor_bits(&table.gradient().expect("table gradient remains"));

    let invalid_target_before = tensor_bits(&logits.gradient().expect("logit gradient exists"));
    let invalid_target = logits
        .indexed_mean_nll(1, &[0, 2])
        .expect_err("class two is outside a two-class row");
    let invalid_target_gradient_unchanged =
        invalid_target_before == tensor_bits(&logits.gradient().expect("logit gradient remains"));

    let empty_targets_before =
        tensor_bits(&empty_logits.gradient().expect("empty gradient exists"));
    let empty_targets = empty_logits
        .indexed_mean_nll(1, &[])
        .expect_err("a mean over no targets is undefined");
    let empty_targets_gradient_unchanged = empty_targets_before
        == tensor_bits(&empty_logits.gradient().expect("empty gradient remains"));

    let exp_overflow_before = tensor_bits(&overflow.gradient().expect("scalar gradient exists"));
    let exp_overflow = overflow
        .exp()
        .expect_err("the finite-forward invariant rejects positive infinity");
    let exp_overflow_gradient_unchanged =
        exp_overflow_before == tensor_bits(&overflow.gradient().expect("scalar gradient remains"));
    let gradients_unchanged = invalid_id_gradient_unchanged
        && invalid_target_gradient_unchanged
        && empty_targets_gradient_unchanged
        && exp_overflow_gradient_unchanged;

    Ok(ModelErrorExample {
        invalid_id,
        invalid_target,
        empty_targets,
        exp_overflow,
        invalid_id_gradient_unchanged,
        invalid_target_gradient_unchanged,
        empty_targets_gradient_unchanged,
        exp_overflow_gradient_unchanged,
        gradients_unchanged,
    })
}

The executable prints the worked values, the three scalar derivatives, the sampled central-difference result, and the aggregate error-preservation result:

Run the complete worked operation chain and print its forward, reverse, numerical, and error results rust/demos/ch16-model-autodiff-ops/src/main.rs#learner-model-vjp-output

    let example = frozen_model_example()?;
    let probes = scalar_probes()?;
    let gradchecks = model_vjp_gradchecks()?;
    let errors = model_error_example()?;

    println!("embeddings: {}", tensor_text(&example.embeddings));
    println!("token IDs: {:?}", example.token_ids);
    println!("gather rows: {}", tensor_text(&example.gathered));
    println!("projection weights: {}", tensor_text(&example.weights));
    println!(
        "projection preactivations: {}",
        tensor_text(&example.projection_preactivations)
    );
    println!("SiLU: {}", tensor_text(&example.activated));
    println!(
        "log-softmax axis=1: {}",
        tensor_text(&example.log_probabilities)
    );
    println!("targets: {:?}", example.targets);
    println!("indexed mean NLL: {}", tensor_text(&example.loss));
    println!(
        "target-logit gradient: {}",
        tensor_text(&example.loss_input_gradient)
    );
    println!(
        "through SiLU: {}",
        tensor_text(&example.matmul_output_gradient)
    );
    println!(
        "matmul left gradient: {}",
        tensor_text(&example.gathered_gradient)
    );
    println!(
        "embedding scatter-add: {}",
        tensor_text(&example.embedding_gradient)
    );
    println!(
        "matmul right gradient: {}",
        tensor_text(&example.weight_gradient)
    );
    println!(
        "scalar probes: exp(0)->({}, {}) | log(1)->({}, {}) | silu(0)->({}, {})",
        fixed(probes[0].output),
        fixed(probes[0].gradient),
        fixed(probes[1].output),
        fixed(probes[1].gradient),
        fixed(probes[2].output),
        fixed(probes[2].gradient),
    );
    println!(
        "gradcheck: {}; pass={}",
        gradchecks
            .checks
            .iter()
            .map(|check| check.operation)
            .collect::<Vec<_>>()
            .join(" | "),
        gradchecks.passed
    );
    println!(
        "typed errors: invalid-id | invalid-target | empty-targets | exp-overflow; gradients unchanged={}",
        errors.gradients_unchanged
    );
    println!("chapter 17 handoff: initialize trainable values reproducibly");

embeddings: shape=3x2 values=2.000000000000,2.000000000000,1.000000000000,-1.000000000000,-1.000000000000,1.000000000000
token IDs: [1, 1, 1, 2]
gather rows: shape=4x2 values=1.000000000000,-1.000000000000,1.000000000000,-1.000000000000,1.000000000000,-1.000000000000,-1.000000000000,1.000000000000
projection weights: shape=2x2 values=1.000000000000,-1.000000000000,1.000000000000,-1.000000000000
projection preactivations: shape=4x2 values=0.000000000000,0.000000000000,0.000000000000,0.000000000000,0.000000000000,0.000000000000,0.000000000000,0.000000000000
SiLU: shape=4x2 values=0.000000000000,0.000000000000,0.000000000000,0.000000000000,0.000000000000,0.000000000000,0.000000000000,0.000000000000
log-softmax axis=1: shape=4x2 values=-0.693147180560,-0.693147180560,-0.693147180560,-0.693147180560,-0.693147180560,-0.693147180560,-0.693147180560,-0.693147180560
targets: [0, 0, 0, 1]
indexed mean NLL: shape=scalar values=0.693147180560
target-logit gradient: shape=4x2 values=-0.125000000000,0.125000000000,-0.125000000000,0.125000000000,-0.125000000000,0.125000000000,0.125000000000,-0.125000000000
through SiLU: shape=4x2 values=-0.062500000000,0.062500000000,-0.062500000000,0.062500000000,-0.062500000000,0.062500000000,0.062500000000,-0.062500000000
matmul left gradient: shape=4x2 values=-0.125000000000,-0.125000000000,-0.125000000000,-0.125000000000,-0.125000000000,-0.125000000000,0.125000000000,0.125000000000
embedding scatter-add: shape=3x2 values=0.000000000000,0.000000000000,-0.375000000000,-0.375000000000,0.125000000000,0.125000000000
matmul right gradient: shape=2x2 values=-0.250000000000,0.250000000000,0.250000000000,-0.250000000000
scalar probes: exp(0)->(1.000000000000, 1.000000000000) | log(1)->(0.000000000000, 1.000000000000) | silu(0)->(0.000000000000, 0.500000000000)
gradcheck: matmul-left | matmul-right | gather_rows | exp | log | silu | log_softmax | indexed_mean_nll; pass=true
typed errors: invalid-id | invalid-target | empty-targets | exp-overflow; gradients unchanged=true
chapter 17 handoff: initialize trainable values reproducibly

Follow four occurrence gradients into three parameter rows

The figure separates the compact forward chain from reverse accumulation. The fork after SiLU is explicit: log-softmax and combined mean NLL are separate calls over the same activated logits, and each retains the probabilities from its own checked forward call. Target rows expose the negative and positive derivatives, their zero class sum, and the exact parent-gradient shapes through SiLU and both matrix operands.

The destination-row boxes contain their own occurrence contributions. Row 1 contains positions 0, 1, and 2 before their sum; row 2 contains position 3; row 0 is visibly unused. That arrangement makes it possible to see why four materialized output rows produce only three parent-table gradient rows.

Follow a repeated token ID from lookup to loss and back

Inspect the compact forward chain, signed target gradients with zero class sums, both matrix VJPs, and four occurrence contributions grouped inside three destination embedding rows.

Token IDs: [1, 1, 1, 2]
Target classes: [0, 0, 0, 1]
Mean token loss: 0.693147180560
Repeated selector: $i=1,\;n=3$

Trace the compact operation chain forward

After SiLU, the graph forks into two separate operation calls over the same logits. Log-softmax returns log-probabilities; combined mean NLL computes and returns one scalar mean NLL from the logits and target classes. The calls do not share one forward result. This exercise is not the final decoder layout.

Row gather

Forward step
0

Inputs
embedding table + token IDs

Differentiable tensor shape

$[3,2]$

Output shape
$[4,2]$

Values
[1.000000000000, -1.000000000000, 1.000000000000, -1.000000000000, 1.000000000000, -1.000000000000, -1.000000000000, 1.000000000000]
Matrix product

Forward step
1

Inputs
gathered feature rows + projection weights

Differentiable tensor shape

$[4,2]$ $[2,2]$

Output shape
$[4,2]$

Values
[0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000]
SiLU

Forward step
2

Inputs
projection preactivations

Differentiable tensor shape

$[4,2]$

Output shape
$[4,2]$

Values
[0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000]
Log-softmax

Forward step
3

Inputs
activated values used as loss logits

Differentiable tensor shape

$[4,2]$

Output shape
$[4,2]$

Values
[-0.693147180560, -0.693147180560, -0.693147180560, -0.693147180560, -0.693147180560, -0.693147180560, -0.693147180560, -0.693147180560]
Indexed mean NLL

Forward step
4

Inputs
activated values used as loss logits + target classes

Differentiable tensor shape

$[4,2]$

Output shape
$[]$

Values
[0.693147180560]

Reverse target selection and the projection

For every flat position in this equal-logit example, begin with the two saved class probabilities, each equal to one half. Subtract one from the target-class component, leave the competing component unchanged, then divide both components by the four positions in the mean.

Loss gradients for the four flat token positions
Flat position	Token ID	Target class	Gradient	Correct-target sign	Other-class sign
0	1	0	`[-0.125000000000, 0.125000000000]`	negative — selected target	positive
1	1	0	`[-0.125000000000, 0.125000000000]`	negative — selected target	positive
2	1	0	`[-0.125000000000, 0.125000000000]`	negative — selected target	positive
3	2	1	`[0.125000000000, -0.125000000000]`	negative — selected target	positive

Transpose the opposite matrix axes, multiply, and reduce broadcast batches to each parent shape.

SiLU

Operand: single input
Parent tensor: projection preactivations
Parent-gradient shape: $[4,2]$
Gradient: [-0.062500000000, 0.062500000000, -0.062500000000, 0.062500000000, -0.062500000000, 0.062500000000, 0.062500000000, -0.062500000000]

Matrix product

Operand: left matrix
Parent tensor: gathered feature rows
Parent-gradient shape: $[4,2]$
Gradient: [-0.125000000000, -0.125000000000, -0.125000000000, -0.125000000000, -0.125000000000, -0.125000000000, 0.125000000000, 0.125000000000]

Matrix product

Operand: right matrix
Parent tensor: projection weights
Parent-gradient shape: $[2,2]$
Gradient: [-0.250000000000, 0.250000000000, 0.250000000000, -0.250000000000]

Add each occurrence to its destination row

Every gathered output row owns its values. In reverse, each occurrence keeps its own adjoint and the gather VJP adds it to the parent-table row named by that token ID.

unused row

Destination row: 0
Flat positions: unused row
Occurrences: 0
Gradient: [0.000000000000, 0.000000000000]

row with repeated contributions

Destination row: 1
Flat positions: 0, 1, 2
Occurrences: 3
Gradient: [-0.375000000000, -0.375000000000]

repeated-ID occurrence

Flat position
0

Occurrence contribution
[-0.125000000000, -0.125000000000]
repeated-ID occurrence

Flat position
1

Occurrence contribution
[-0.125000000000, -0.125000000000]
repeated-ID occurrence

Flat position
2

Occurrence contribution
[-0.125000000000, -0.125000000000]

row with one contribution

Destination row: 2
Flat positions: 3
Occurrences: 1
Gradient: [0.125000000000, 0.125000000000]

single-ID occurrence

Flat position
3

Occurrence contribution
[0.125000000000, 0.125000000000]

Predict before running Rust

Map flat positions 0 through 3 to $(b,t)$ for $B=1$ and $T=4$ .
Write all four gathered rows, projection-preactivation rows, loss-logit rows, log-probability rows, and the scalar mean loss.
Predict why the correct-target gradient is negative, the competing gradient is positive, and each two-class row sums to zero.
Apply the SiLU derivative at zero. What is the gradient entering matmul at each position?
Predict the shapes of both matmul parent gradients and calculate $dW$ .
List the three contributions to embedding row 1 before adding them. What reaches rows 0 and 2?
Predict both value and local derivative for $\exp(0)$ , $\ln(1)$ , and $\operatorname{SiLU}(0)$ .
Explain why logits near $\pm1000$ require maximum-shifted normalization and why each autodiff probability operation saves the probabilities emitted by its own checked forward call. Does “one call” mean one read of each logit? In what sense are the saved f64 values identical, and do the two probability branches share a result?
Which row-gather facts are checked at the public entry, and which work may the validated plan omit? At what stage are an invalid target class and an empty target set detected?
Misconception check: should row 1 be divided by three after its three occurrence gradients are added?

Check the ten model-operation predictions

The map is $0\mapsto(0,0)$ , $1\mapsto(0,1)$ , $2\mapsto(0,2)$ , and $3\mapsto(0,3)$ .
Gather gives [1,-1] three times and [-1,1] once. Every projection preactivation, SiLU output, and loss-logit row is [0,0]; every log-probability row is $[-\ln 2,-\ln 2]$ ; the mean loss is $\ln 2$ .
Increasing the selected target logit reduces NLL, so its derivative is negative. Increasing a competitor raises its probability at the target’s expense, so its derivative is positive. Softmax probabilities sum to one and the one-hot target also sums to one, leaving class-row gradient sum zero.
SiLU has derivative $1/2$ at zero. The first three rows become $[-1/16,1/16]$ ; the last becomes $[1/16,-1/16]$ .
The left-parent gradient has gathered shape [4,2]; the right-parent gradient has weight shape [2,2]. The value is $dW=[[-1/4,1/4],[1/4,-1/4]]$ .
Positions 0, 1, and 2 each contribute $[-1/8,-1/8]$ to row 1, so it gets $[-3/8,-3/8]$ . Row 2 gets $[1/8,1/8]$ ; unused row 0 gets [0,0].
$\exp(0)=1$ with derivative $1$ ; $\ln(1)=0$ with derivative $1$ ; $\operatorname{SiLU}(0)=0$ with derivative $1/2$ .
Maximum shifting prevents avoidable overflow without changing the distribution. Each log-softmax or indexed-mean-NLL operation emits its primary result and the probability values needed by its VJP in one checked call. The saved tensor contains the same emitted f64 values, bit for bit; this is identity of stored floating-point values, not exact real arithmetic. Backward uses those values without another softmax call. “One call” does not mean one logit read: the preliminary scan checks finiteness, RowStats uses separate class scans for the maximum and shifted-exponential sum, and the emitter writes the probability output. The lesson’s two branches are separate operations with separate calls and saved tensors; they share only the input logits.
The public row-gather entry first checks operand availability, table rank, the logical selector shape, the selector count, and then IDs in flat order. Its validated plan lets row copying and scatter-add reuse those facts without rescanning the raw request. The indexed-loss entry validates the class axis, checks that the target count equals the number of class-axis groups, rejects zero groups, and then checks targets in flat group order before computing the loss or saved probabilities.
No. Every contribution already includes the loss’s $1/4$ factor. Gather reversal sums the three contributions exactly once and does not apply another mean.

Initialize the values these gradients will train

The cumulative implementation can now differentiate a compact chain from selected embedding rows through a projection, nonlinearity, and stable mean token loss. This is an operation-level example, not the final decoder architecture: later feed-forward blocks use SiLU internally, and a separate vocabulary projection produces the decoder’s loss logits.

Later chapters will package these operations into embedding, linear, normalization, SwiGLU, and attention layers. Correct gradients still do not choose useful parameter values. The next problem is choosing initial values that break symmetry without making signals vanish or explode, so Chapter 17 adds deterministic, scale-aware initialization without adding another VJP.