15 · Content revision 9

Reverse tensor operations with edge-local VJPs

Build a Rust tensor autodiff tape, reverse shape transformations, broadcasts, and reductions with edge-local VJPs, and verify gradients for LLM training.

Predict one shape-changing tensor graph

Begin with two parameter leaves:

x shape [2,3]       bias shape [3]
[[1, 2, 3],         [1, -1, 0]
 [4, 5, 6]]

Reshape x to [3,2], then transpose axes 0 and 1. Reshape preserves logical row-major order, and transpose changes which coordinate reads each value:

reshape [3,2]       transpose [2,3]
[[1, 2],            [[1, 3, 5],
 [3, 4],             [2, 4, 6]]
 [5, 6]]

Explicitly broadcast bias across the leading axis, add it, square by using the same add result as both multiply operands, and mean axis 1:

broadcast bias      add                multiply           mean axis 1
[[ 1,-1,0],         [[2,2,5],          [[4,4,25],         [11,18]
 [ 1,-1,0]]          [3,3,6]]           [9,9,36]]

The output is not scalar. Select the scalar objective $L=3y_0+6y_1$ ; its output cotangent is therefore $\bar{y}=[3,6]$ . The mean VJP divides each seed value by the saved axis extent three, then repeats that quotient across the corresponding three-element input row. Thus [3,6] becomes $[3/3,3/3,3/3,6/3,6/3,6/3]=[1,1,1,2,2,2]$ at multiply. Both multiply operands are the same node. Each edge contributes upstream times the other operand, so their sum is:

\bar{\mathrm{add}}=[4,4,10,12,12,24]

Add’s VJP copies that completed adjoint once to the transposed-x input and once to the broadcast-bias input; it does not divide the adjoint between them. Broadcast reversal then adds the two reused rows coordinate by coordinate: $[4+12,4+12,10+24]$ , giving $\bar{\mathrm{bias}}=[16,16,34]$ . Transpose swaps the same axes and reshape restores $x$ ‘s original shape, giving $\bar{x}=[4,12,4,12,10,24]$ .

Contrast one fixed-shape backward calculation with the reusable eight-node tensor tape rust/demos/ch15-tensor-autodiff-core/src/lib.rs#shared-tensor-vjp-fixture

/// Computes the worked graph with bounded arrays and a handwritten backward pass.
///
/// This is intentionally a contrast, not a second autodiff implementation: every
/// loop bound and shape is fixed to the eight-node Chapter 15 fixture.
pub fn handwritten_fixed_shape_baseline() -> HandwrittenBaseline {
    let mut reshaped = [0.0; 6];
    reshaped.copy_from_slice(&X_VALUES);

    let mut transposed = [0.0; 6];
    for row in 0..RESHAPED_SHAPE[0] {
        for column in 0..RESHAPED_SHAPE[1] {
            transposed[column * X_SHAPE[1] + row] = reshaped[row * RESHAPED_SHAPE[1] + column];
        }
    }

    let mut broadcast = [0.0; 6];
    let mut added = [0.0; 6];
    let mut squared = [0.0; 6];
    for row in 0..X_SHAPE[0] {
        for (column, &bias_value) in BIAS_VALUES.iter().enumerate() {
            let offset = row * X_SHAPE[1] + column;
            broadcast[offset] = bias_value;
            added[offset] = transposed[offset] + broadcast[offset];
            squared[offset] = added[offset] * added[offset];
        }
    }

    let mut output = [0.0; 2];
    for row in 0..X_SHAPE[0] {
        output[row] = squared[row * X_SHAPE[1]..(row + 1) * X_SHAPE[1]]
            .iter()
            .sum::<f64>()
            / X_SHAPE[1] as f64;
    }

    let mut added_adjoint = [0.0; 6];
    for (row, &seed) in SEED_VALUES.iter().enumerate() {
        let mean_adjoint = seed / X_SHAPE[1] as f64;
        for column in 0..X_SHAPE[1] {
            let offset = row * X_SHAPE[1] + column;
            // The repeated multiply has two ordered operand contributions.
            added_adjoint[offset] += mean_adjoint * added[offset];
            added_adjoint[offset] += mean_adjoint * added[offset];
        }
    }

    let mut bias_gradient = [0.0; 3];
    for row in 0..X_SHAPE[0] {
        for column in 0..X_SHAPE[1] {
            bias_gradient[column] += added_adjoint[row * X_SHAPE[1] + column];
        }
    }

    let mut reshaped_adjoint = [0.0; 6];
    for row in 0..RESHAPED_SHAPE[0] {
        for column in 0..RESHAPED_SHAPE[1] {
            reshaped_adjoint[row * RESHAPED_SHAPE[1] + column] =
                added_adjoint[column * X_SHAPE[1] + row];
        }
    }
    let x_gradient = reshaped_adjoint;

    HandwrittenBaseline {
        output,
        x_gradient,
        bias_gradient,
    }
}

fn frozen_x() -> Tensor {
    tensor(&X_SHAPE, &X_VALUES)
}

fn frozen_bias() -> Tensor {
    tensor(&BIAS_SHAPE, &BIAS_VALUES)
}

fn frozen_seed() -> Tensor {
    tensor(&OUTPUT_SHAPE, &SEED_VALUES)
}

/// The accumulated gradients stored by the two parameter leaves.
#[derive(Clone, Debug, PartialEq)]
pub struct ParameterGradients {
    pub x: Tensor,
    pub bias: Tensor,
}

/// Forward nodes, reverse evidence, and all retained/released parameter states.
#[derive(Clone, Debug)]
pub struct FrozenTensorExample {
    pub nodes: Vec<TensorValue>,
    pub seed: Tensor,
    pub baseline: HandwrittenBaseline,
    pub first_pass: TensorBackwardPass,
    pub first: ParameterGradients,
    pub repeated: ParameterGradients,
    pub zeroed: ParameterGradients,
    pub after_zero_and_release: ParameterGradients,
    pub released_error: TensorAutodiffError,
    pub released_gradients_unchanged: bool,
}

fn parameter_snapshot(x: &TensorValue, bias: &TensorValue) -> ParameterGradients {
    ParameterGradients {
        x: x.gradient_snapshot().expect("x is a parameter"),
        bias: bias.gradient_snapshot().expect("bias is a parameter"),
    }
}

fn tensor_bits(value: &Tensor) -> Vec<u64> {
    value.as_slice().iter().map(|item| item.to_bits()).collect()
}

fn optional_tensor_bits(value: Option<Ref<'_, Tensor>>) -> Option<Vec<u64>> {
    value.as_ref().map(|value| tensor_bits(value))
}

/// Builds the same expression through reusable operation-level VJPs and checks
/// its result against `handwritten_fixed_shape_baseline` before returning it.
pub fn frozen_tensor_example() -> Result<FrozenTensorExample, TensorAutodiffError> {
    let x = TensorValue::parameter(frozen_x())?;
    let reshaped = x.reshape(&RESHAPED_SHAPE)?;
    let transposed = reshaped.transpose(0, 1)?;
    let bias = TensorValue::parameter(frozen_bias())?;
    let broadcast = bias.broadcast_to(&X_SHAPE)?;
    let added = transposed.add(&broadcast)?;
    let multiplied = added.mul(&added)?;
    let output = multiplied.mean_axis(1, false)?;
    let seed = frozen_seed();
    let baseline = handwritten_fixed_shape_baseline();

    assert_eq!(output.value().as_slice(), baseline.output);
    let first_pass = output.backward_with_seed_and_trace(&seed.view(), GraphRetention::Retain)?;
    let first = parameter_snapshot(&x, &bias);
    assert_eq!(first.x.as_slice(), baseline.x_gradient);
    assert_eq!(first.bias.as_slice(), baseline.bias_gradient);

    output.backward_with_seed(&seed.view(), GraphRetention::Retain)?;
    let repeated = parameter_snapshot(&x, &bias);
    x.zero_grad()?;
    bias.zero_grad()?;
    let zeroed = parameter_snapshot(&x, &bias);
    output.backward_with_seed(&seed.view(), GraphRetention::Release)?;
    let after_zero_and_release = parameter_snapshot(&x, &bias);
    let released_before = (
        optional_tensor_bits(x.gradient()),
        optional_tensor_bits(bias.gradient()),
    );
    let released_error = output
        .backward_with_seed(&seed.view(), GraphRetention::Retain)
        .expect_err("a released operation result must reject another pass");
    let released_gradients_unchanged = released_before
        == (
            optional_tensor_bits(x.gradient()),
            optional_tensor_bits(bias.gradient()),
        );

    Ok(FrozenTensorExample {
        nodes: vec![
            x, reshaped, transposed, bias, broadcast, added, multiplied, output,
        ],
        seed,
        baseline,
        first_pass,
        first,
        repeated,
        zeroed,
        after_zero_and_release,
        released_error,
        released_gradients_unchanged,
    })
}

Apply one edge-local VJP instead of building a Jacobian

The central reverse rule is:

\bar{p(e)}\mathrel{+}=J_e^\top\bar{c(e)},\qquad e\in E

Each reachable operand-use edge $e\in E$ connects one parent tensor $p(e)$ to the consumer result $c(e)$ . The local map $J_e$ differentiates $c(e)$ with respect to that one operand slot while the operation’s other slots stay fixed. Reverse mode already has $\bar{c(e)}$ , the selected objective’s sensitivity to each coordinate of the consumer. Applying $J_e^\top$ maps that information back to the exact shape of $p(e)$ .

The VJP applies this map directly without creating a matrix with one row per consumer coordinate and one column per parent coordinate. The $\mathrel{+}=$ matters because another branch—or another operand slot of the same operation—may reach the same parent. Thus $\mathrm{add}\odot\mathrm{add}$ contributes through two distinct edges even though both edges refer to one node.

Name the tensors, map, and adjoints

Symbol	Operational meaning
$E$	The reachable set of recorded operand-use edges.
$e$	One distinct operand occurrence at one consumer operation.
$p(e)$	The parent tensor supplied through edge $e$ .
$c(e)$	The result of the operation that consumes edge $e$ .
$J_e$	The conceptual Jacobian for that operand slot, with the operation’s other slots fixed.
$\bar{c(e)}$	The fresh upstream adjoint with exactly the shape of $c(e)$ .
$\bar{p(e)}$	The fresh parent-adjoint accumulator with exactly the shape of $p(e)$ .
$\mathrel{+}=$	Add this edge’s contribution instead of overwriting other paths.

The superscript transpose belongs to the conceptual slot-local Jacobian. The graph’s forward transpose(0,1) is a separate tensor operation whose VJP happens to swap those saved axes again.

From explicit next-word updates to reusable tensor pullbacks

Bengio et al.’s neural language model has millions of parameters and an explicit forward phase followed by network-specific backward/update equations. Those equations make gradient flow from a next-word loss back to the model parameters inspectable. But a separate graph node for every scalar value—or a separately handwritten backward calculation for every tensor expression—becomes unwieldy in deep models with many repeated blocks that change tensor shapes.

The earlier neural-language-model checkpoint is Bengio et al., A Neural Probabilistic Language Model. Bengio et al. describe a neural next-word model with millions of parameters and publish an explicit forward phase followed by backward/update equations for output, hidden, and learned word-feature gradients.

The fixed-shape calculation in the worked example makes that boundary concrete. It can produce the right gradients for exactly one known expression, but its backward arithmetic is not reusable when a new shape-changing operation is inserted. It is an illustrative contrast, not code attributed to Bengio et al.

Abadi et al. represent computation as operation vertices joined by tensor-valued edges and describe automatic differentiation that finds every backward path from a loss to parameters and sums the paths’ partial gradients. Vaswani et al. then train repeated Transformer attention and feed-forward tensor blocks, while Radford et al. scale autoregressive Transformer language models to deeper and wider stacks.

The operation-graph checkpoint is Abadi et al., TensorFlow: A System for Large-Scale Machine Learning. Abadi et al. define graph vertices as operations and edge values as tensors, then describe a differentiation library that derives backpropagation for layer-and-loss compositions by finding backward paths to parameters and summing each path’s partial-gradient contribution.

The Transformer-training checkpoint is Vaswani et al., Attention Is All You Need. Vaswani et al. build the Transformer from repeated attention and position-wise feed-forward sublayers and train base models for 100,000 steps and big models for 300,000 steps with Adam.

The scaled autoregressive-LM checkpoint is Radford et al., Language Models are Unsupervised Multitask Learners. Radford et al. use Transformer-based autoregressive language models and report four sizes spanning 12 to 48 layers and 117 million to 1.542 billion parameters.

This chapter records one local vector-Jacobian product for each operand use, restores every contribution to its parent’s exact shape through reshape, transpose, broadcast, sum, and mean rules, and checks those rules numerically before model-specific derivatives are added. Ordinary inference does not run the reverse tape; training uses it to carry loss sensitivity back through repeated tensor blocks.

Abadi et al. support the operation-and-tensor graph and summed-path claims, but their symbolic graph does not prescribe this implementation’s owned eager tape. Vaswani et al.’s Transformer paper establishes repeated attention and feed-forward blocks during training, while Radford et al.’s GPT-2 report establishes deeper and wider autoregressive Transformer models. Neither source prescribes these VJP rules or lifecycle choices.

Own tensor primals and save only local context

TensorValue::parameter and TensorValue::constant accept finite contiguous tensors. Each tape node owns its primal tensor, including every parameter, constant, and operation result. Calling value() lends a temporary read-only guard to that node-owned primal; it does not copy the tensor. Calling gradient() similarly lends a read-only guard when a parameter leaf has a stored gradient. Constants and operation nodes do not store parameter gradients.

A read guard is temporary and must be dropped before code can mutate the same storage. If a stored-gradient guard is still active, zero_grad or a reverse pass returns GradientBorrowed instead of panicking or committing only part of a pass. When data must remain independently owned after the read ends, value_snapshot() clones the primal and gradient_snapshot() clones the stored gradient when one exists.

A tensor snapshot contains data but is not a tape node. Cloning a TensorValue keeps the original node identity. detach() does something different from both: it snapshots the primal, then creates a new untracked TensorValue leaf with no operand edge. Later changes to either node’s storage cannot change the other leaf’s independently owned primal.

Each node also carries an internal revision for its primal tensor. A later training step can write an updated parameter value into the same node, preserving node identity while advancing that primal revision. Node identity alone therefore does not prove that an old retained graph still describes the current parameter values.

Primal revisions are runtime tape-validity metadata: backward uses them only to decide whether saved graph context still belongs to the current parameter values. They are not optimizer step numbers and are not serialized in model checkpoints.

Own finite tensor primals while distinguishing guarded reads, independent snapshots, identity, detach, and release rust/crates/llm-from-scratch/src/autograd/tensor_core.rs#tensor-tape-values

#[derive(Clone)]
struct ParentEdge {
    parent: TensorValue,
    parent_value_revision: u64,
    saved: TensorSavedContext,
}

impl ParentEdge {
    fn capture(parent: &TensorValue, saved: TensorSavedContext) -> Self {
        Self {
            parent: parent.clone(),
            parent_value_revision: parent.value_revision(),
            saved,
        }
    }
}

struct NodeState {
    parents: Vec<ParentEdge>,
    parameter_gradient: Option<Tensor>,
    released: bool,
}

struct Node {
    value: RefCell<Tensor>,
    value_revision: Cell<u64>,
    operation: TensorOperation,
    tracked: bool,
    state: RefCell<NodeState>,
}

/// One owned tensor value and its operation-level reverse-mode tape.
#[derive(Clone)]
pub struct TensorValue {
    node: Rc<Node>,
}

impl fmt::Debug for TensorValue {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("TensorValue")
            .field("shape", &self.shape())
            .field("operation", &self.operation())
            .field("tracks_gradient", &self.tracks_gradient())
            .field("gradient", &self.gradient())
            .field("released", &self.is_released())
            .finish()
    }
}

/// One node in the deterministic parent-first topology of a backward pass.
#[derive(Clone, Debug, PartialEq)]
pub struct TensorBackwardNode {
    pub topology_index: usize,
    pub operation: TensorOperation,
    pub shape: Vec<usize>,
    pub tracked: bool,
    pub parameter: bool,
    pub pass_adjoint: Option<Tensor>,
    pub accumulated_gradient: Option<Tensor>,
}

/// One ordered operand edge visited during tensor reverse traversal.
#[derive(Clone, Debug, PartialEq)]
pub struct TensorBackwardEdge {
    pub reverse_index: usize,
    pub child: usize,
    pub parent: usize,
    pub operand: usize,
    pub saved: TensorSavedContext,
    pub upstream: Tensor,
    pub contribution: Tensor,
    pub parent_tracked: bool,
    pub parent_adjoint_before: Option<Tensor>,
    pub parent_adjoint_after: Option<Tensor>,
}

/// The trace returned by one explicitly observed, successfully committed pass.
#[derive(Clone, Debug, PartialEq)]
pub struct TensorBackwardPass {
    pub seed: Tensor,
    pub retention: GraphRetention,
    pub nodes: Vec<TensorBackwardNode>,
    pub edges: Vec<TensorBackwardEdge>,
}

impl TensorValue {
    /// Creates a finite leaf parameter initialized with an exact-shape zero gradient.
    pub fn parameter(value: Tensor) -> Result<Self, TensorAutodiffError> {
        check_finite_leaf(&value, TensorOperation::Parameter)?;
        let gradient = zeros(value.shape())?;
        Ok(Self::new_node(
            value,
            TensorOperation::Parameter,
            Vec::new(),
            true,
            Some(gradient),
        ))
    }

    /// Creates a finite untracked tensor leaf.
    pub fn constant(value: Tensor) -> Result<Self, TensorAutodiffError> {
        check_finite_leaf(&value, TensorOperation::Constant)?;
        Ok(Self::new_node(
            value,
            TensorOperation::Constant,
            Vec::new(),
            false,
            None,
        ))
    }

    fn new_node(
        value: Tensor,
        operation: TensorOperation,
        parents: Vec<ParentEdge>,
        tracked: bool,
        parameter_gradient: Option<Tensor>,
    ) -> Self {
        Self {
            node: Rc::new(Node {
                value: RefCell::new(value),
                value_revision: Cell::new(0),
                operation,
                tracked,
                state: RefCell::new(NodeState {
                    parents,
                    parameter_gradient,
                    released: false,
                }),
            }),
        }
    }

    fn operation_node(
        value: Tensor,
        operation: TensorOperation,
        mut parents: Vec<ParentEdge>,
    ) -> Result<Self, TensorAutodiffError> {
        check_finite_forward(&value, operation)?;
        let tracked = !no_grad_active() && parents.iter().any(|edge| edge.parent.tracks_gradient());
        if no_grad_active() {
            parents.clear();
        }
        Ok(Self::new_node(value, operation, parents, tracked, None))
    }

    /// Builds one checked model-operation node without exposing tape internals.
    pub(crate) fn model_operation<const N: usize>(
        operation: TensorOperation,
        operands: [&Self; N],
        forward: impl FnOnce(
            [&Tensor; N],
        ) -> Result<(Tensor, [ModelSavedContext; N]), TensorAutodiffError>,
    ) -> Result<Self, TensorAutodiffError> {
        ensure_operands_available(operation, &operands)?;
        let primals: [Ref<'_, Tensor>; N] = std::array::from_fn(|index| operands[index].value());
        let primal_refs: [&Tensor; N] = std::array::from_fn(|index| &*primals[index]);
        let (value, contexts) = forward(primal_refs)?;
        let parents = operands
            .into_iter()
            .zip(contexts)
            .map(|(parent, context)| {
                ParentEdge::capture(parent, TensorSavedContext::Model(context))
            })
            .collect();
        Self::operation_node(value, operation, parents)
    }

    /// Borrows the node-owned primal tensor.
    pub fn value(&self) -> Ref<'_, Tensor> {
        self.node.value.borrow()
    }

    /// Copies the node-owned primal into an independent tensor snapshot.
    pub fn value_snapshot(&self) -> Tensor {
        self.node.value.borrow().clone()
    }

    /// Copies the primal shape.
    pub fn shape(&self) -> Vec<usize> {
        self.node.value.borrow().shape().to_vec()
    }

    pub fn operation(&self) -> TensorOperation {
        self.node.operation
    }

    pub fn tracks_gradient(&self) -> bool {
        self.node.tracked
    }

    pub fn is_parameter(&self) -> bool {
        self.operation() == TensorOperation::Parameter
    }

    pub fn is_released(&self) -> bool {
        self.node.state.borrow().released
    }

    /// Borrows the accumulated gradient stored only by a parameter leaf.
    pub fn gradient(&self) -> Option<Ref<'_, Tensor>> {
        Ref::filter_map(self.node.state.borrow(), |state| {
            state.parameter_gradient.as_ref()
        })
        .ok()
    }

    /// Copies the accumulated parameter gradient into an independent snapshot.
    pub fn gradient_snapshot(&self) -> Option<Tensor> {
        self.gradient().as_deref().cloned()
    }

    /// Returns whether two handles refer to the same tape node.
    pub fn is_same_node(&self, other: &Self) -> bool {
        Rc::ptr_eq(&self.node, &other.node)
    }

    /// Copies the primal into a new untracked leaf and severs all parent edges.
    pub fn detach(&self) -> Self {
        Self::new_node(
            self.value_snapshot(),
            TensorOperation::Detached,
            Vec::new(),
            false,
            None,
        )
    }
}

Forward operations materialize finite contiguous results. Add saves the two operand shapes. Each multiply edge saves its parent shape and the other operand’s primal because its local derivative needs that value. Reshape saves its input shape, transpose saves its two axes, broadcast saves its source shape, and reductions save their axis, retained-dimension choice, and input extent. These are owned operation values, not borrowed zero-copy views.

When a forward operation creates an operand-use edge, that edge captures the parent’s current primal revision. The revision binds the edge and its saved VJP context to the exact parent value used in this forward calculation. If an in-place parameter update later advances the revision, the old edge still reaches the same node, but the retained context belongs to the earlier value.

Record checked tensor operations with ordered edges and operation-specific VJP context rust/crates/llm-from-scratch/src/autograd/tensor_core.rs#tensor-forward-operations

impl TensorValue {
    /// Adds two tensors using trailing-axis broadcasting.
    pub fn add(&self, other: &Self) -> Result<Self, TensorAutodiffError> {
        ensure_operands_available(TensorOperation::Add, &[self, other])?;
        let left = self.value();
        let right = other.value();
        let value = map_binary(&left.view(), &right.view(), |a, b| a + b)?;
        let output_shape = value.shape().to_vec();
        let parents = vec![
            ParentEdge::capture(self, broadcast_context(left.shape(), &output_shape)),
            ParentEdge::capture(other, broadcast_context(right.shape(), &output_shape)),
        ];
        Self::operation_node(value, TensorOperation::Add, parents)
    }

    /// Multiplies two tensors and records one ordered edge per operand use.
    pub fn mul(&self, other: &Self) -> Result<Self, TensorAutodiffError> {
        ensure_operands_available(TensorOperation::Multiply, &[self, other])?;
        let left = self.value();
        let right = other.value();
        let value = map_binary(&left.view(), &right.view(), |a, b| a * b)?;
        let output_shape = value.shape().to_vec();
        let parents = vec![
            ParentEdge::capture(
                self,
                multiply_context(left.shape(), &output_shape, Tensor::clone(&right)),
            ),
            ParentEdge::capture(
                other,
                multiply_context(right.shape(), &output_shape, Tensor::clone(&left)),
            ),
        ];
        Self::operation_node(value, TensorOperation::Multiply, parents)
    }

    /// Changes shape without changing row-major element order.
    pub fn reshape(&self, shape: &[usize]) -> Result<Self, TensorAutodiffError> {
        ensure_operands_available(TensorOperation::Reshape, &[self])?;
        let input = self.value();
        let value = input.view().reshape(shape)?.materialize()?;
        let saved = TensorSavedContext::Reshape {
            input_shape: input.shape().to_vec(),
            output_shape: value.shape().to_vec(),
        };
        Self::operation_node(
            value,
            TensorOperation::Reshape,
            vec![ParentEdge::capture(self, saved)],
        )
    }

    /// Swaps two axes and materializes the logical result as owned storage.
    pub fn transpose(
        &self,
        first_axis: usize,
        second_axis: usize,
    ) -> Result<Self, TensorAutodiffError> {
        ensure_operands_available(TensorOperation::Transpose, &[self])?;
        let input = self.value();
        let value = input
            .view()
            .transpose(first_axis, second_axis)?
            .materialize()?;
        let saved = TensorSavedContext::Transpose {
            first_axis,
            second_axis,
            input_shape: input.shape().to_vec(),
            output_shape: value.shape().to_vec(),
        };
        Self::operation_node(
            value,
            TensorOperation::Transpose,
            vec![ParentEdge::capture(self, saved)],
        )
    }

    /// Broadcasts exactly to `shape`; the requested shape may only expand axes.
    pub fn broadcast_to(&self, shape: &[usize]) -> Result<Self, TensorAutodiffError> {
        ensure_operands_available(TensorOperation::Broadcast, &[self])?;
        let input = self.value();
        let inferred = broadcast_shape(input.shape(), shape)?;
        if inferred != shape {
            return Err(TensorAutodiffError::BroadcastTargetMismatch {
                input: input.shape().to_vec(),
                requested: shape.to_vec(),
                inferred,
            });
        }
        let blank = zeros(shape)?;
        let value = map_binary(&input.view(), &blank.view(), |value, _| value)?;
        Self::operation_node(
            value,
            TensorOperation::Broadcast,
            vec![ParentEdge::capture(
                self,
                broadcast_context(input.shape(), shape),
            )],
        )
    }

    /// Sums one axis and records how to expand its exact-shape VJP.
    pub fn sum_axis(&self, axis: usize, keep_dim: bool) -> Result<Self, TensorAutodiffError> {
        self.reduce_axis(axis, keep_dim, false)
    }

    /// Averages one nonempty axis and records the divisor for its VJP.
    pub fn mean_axis(&self, axis: usize, keep_dim: bool) -> Result<Self, TensorAutodiffError> {
        self.reduce_axis(axis, keep_dim, true)
    }

    fn reduce_axis(
        &self,
        axis: usize,
        keep_dim: bool,
        mean: bool,
    ) -> Result<Self, TensorAutodiffError> {
        let operation = if mean {
            TensorOperation::Mean
        } else {
            TensorOperation::Sum
        };
        ensure_operands_available(operation, &[self])?;
        let input = self.value();
        let value = if mean {
            tensor_mean_axis(&input.view(), axis, keep_dim)?
        } else {
            tensor_sum_axis(&input.view(), axis, keep_dim)?
        };
        let divisor = if mean { input.shape()[axis] } else { 1 };
        let saved = TensorSavedContext::Reduction {
            axis,
            keep_dim,
            divisor,
            input_shape: input.shape().to_vec(),
            output_shape: value.shape().to_vec(),
        };
        Self::operation_node(value, operation, vec![ParentEdge::capture(self, saved)])
    }
}

Each local rule returns its parent shape. Add and multiply first form one contribution per operand edge, then reduce any trailing-axis broadcast back to that operand’s saved shape. Reshape restores logical flat values into the saved shape. Transpose swaps the same axes. Explicit broadcast sums missing leading and singleton parent axes. Sum expands its upstream adjoint along the reduced axis; mean performs the same expansion and divides every copied value by the saved nonzero extent. With keep_dim=false, the reduced axis is absent from the upstream shape and is inserted with effective stride zero. With keep_dim=true, the upstream keeps that axis at extent one, but the effective stride is still zero because every input coordinate along the axis reads the same upstream value.

The implementation represents broadcast reversal and reduction expansion with a projected-stride plan. The projection maps every axis of one logical traversal shape to an effective stride in the storage being read or written. An effective stride is the storage-offset contribution of increasing that axis’s coordinate by one while the other coordinates stay fixed. The VJP derives these strides from operation metadata that the forward pass already validated: reused, missing, or restored axes map to stride zero, while ordinary aligned axes keep their row-major strides.

The checked cursor verifies the supplied plan’s rank, recomputes its logical element count, rejects offset arithmetic that would overflow, and, for a nonempty traversal, proves that the largest reachable offset fits the addressed backing slice. An empty traversal emits no offset and performs no read. After that one safety check, the cursor advances in row-major order with one coordinate record per axis, or $O(\mathrm{rank})$ state. It does not rebuild a coordinate vector or call a coordinate-to-offset lookup for every scalar.

Reduction expansion uses effective strides to select upstream source values. For the worked mean, the mean input and multiply result have shape [2,3], while the incoming adjoint has shape [2]. Effective source strides [1,0] retain the upstream row stride and assign zero to the restored mean axis. They produce source offsets [0,0,0,1,1,1]: the first three destination values read seed position zero, and the next three read seed position one. The newly allocated gradient for the mean input is contiguous, so it is written at destination offsets [0,1,2,3,4,5]; each read is divided by three. For the same shape, axis, and keep_dim choice, sum uses the same offsets with divisor one. In general, all non-reduced axes keep the upstream tensor’s row-major strides.

Broadcast reversal instead uses effective strides to select destination accumulators. The broadcast operation’s output—and therefore the incoming adjoint to its VJP—has shape [2,3], while the newly allocated bias gradient has shape [3]. Destination strides [0,1] assign zero to the missing leading axis and retain the bias gradient’s trailing stride. The loop reads incoming values [4,4,10,12,12,24] in flat row-major order while the cursor maps them to destination offsets [0,1,2,0,1,2]. The additions therefore occur as $[4+12,4+12,10+24]$ and produce [16,16,34] in the original order. Every missing leading axis and every aligned parent axis whose extent is one receives effective stride zero; every aligned non-singleton axis retains its parent’s row-major stride.

Traverse checked projected offsets while reversing broadcasts and reductions to exact parent shapes rust/crates/llm-from-scratch/src/autograd/tensor_core.rs#tensor-structural-vjps

fn broadcast_context(input_shape: &[usize], output_shape: &[usize]) -> TensorSavedContext {
    TensorSavedContext::Broadcast {
        input_shape: input_shape.to_vec(),
        output_shape: output_shape.to_vec(),
        reduced_axes: broadcast_reduced_axes(input_shape, output_shape),
    }
}

fn multiply_context(
    input_shape: &[usize],
    output_shape: &[usize],
    other: Tensor,
) -> TensorSavedContext {
    TensorSavedContext::Multiply {
        other,
        input_shape: input_shape.to_vec(),
        output_shape: output_shape.to_vec(),
        reduced_axes: broadcast_reduced_axes(input_shape, output_shape),
    }
}

fn broadcast_reduced_axes(input_shape: &[usize], output_shape: &[usize]) -> Vec<usize> {
    let padding = output_shape.len() - input_shape.len();
    (0..output_shape.len())
        .filter(|&axis| {
            axis < padding || (input_shape[axis - padding] == 1 && output_shape[axis] != 1)
        })
        .collect()
}

fn apply_vjp(upstream: &Tensor, saved: &TensorSavedContext) -> Result<Tensor, TensorAutodiffError> {
    match saved {
        TensorSavedContext::Broadcast {
            input_shape,
            output_shape,
            ..
        } => {
            debug_assert_eq!(upstream.shape(), output_shape);
            unbroadcast(upstream, input_shape)
        }
        TensorSavedContext::Multiply {
            other,
            input_shape,
            output_shape,
            ..
        } => {
            debug_assert_eq!(upstream.shape(), output_shape);
            let product = map_binary(&upstream.view(), &other.view(), |a, b| a * b)?;
            unbroadcast(&product, input_shape)
        }
        TensorSavedContext::Reshape {
            input_shape,
            output_shape,
        } => {
            debug_assert_eq!(upstream.shape(), output_shape);
            Ok(upstream.view().reshape(input_shape)?.materialize()?)
        }
        TensorSavedContext::Transpose {
            first_axis,
            second_axis,
            output_shape,
            ..
        } => {
            debug_assert_eq!(upstream.shape(), output_shape);
            Ok(upstream
                .view()
                .transpose(*first_axis, *second_axis)?
                .materialize()?)
        }
        TensorSavedContext::Reduction {
            axis,
            keep_dim,
            divisor,
            input_shape,
            output_shape,
        } => {
            debug_assert_eq!(upstream.shape(), output_shape);
            expand_reduction(upstream, input_shape, *axis, *keep_dim, *divisor)
        }
        TensorSavedContext::Model(saved) => apply_model_vjp(upstream, saved),
    }
}

fn unbroadcast(upstream: &Tensor, input_shape: &[usize]) -> Result<Tensor, TensorAutodiffError> {
    let mut result = zeros(input_shape)?;
    accumulate_unbroadcast(upstream, &mut result);
    Ok(result)
}

pub(super) fn accumulate_unbroadcast(upstream: &Tensor, result: &mut Tensor) {
    let output_shape = upstream.shape();
    let input_shape = result.shape();
    debug_assert!(output_shape.len() >= input_shape.len());
    let padding = output_shape.len() - input_shape.len();

    let destination_strides = output_shape
        .iter()
        .enumerate()
        .map(|(output_axis, _)| {
            if output_axis < padding || input_shape[output_axis - padding] == 1 {
                0
            } else {
                result.strides()[output_axis - padding]
            }
        })
        .collect::<Vec<_>>();
    let destination_offsets = result
        .view()
        .projected_offsets(output_shape, &destination_strides, upstream.len())
        .expect("a checked broadcast VJP retains a valid destination traversal plan");

    for (&value, destination_offset) in upstream.as_slice().iter().zip(destination_offsets) {
        result.as_mut_slice()[destination_offset] += value;
    }
}

fn expand_reduction(
    upstream: &Tensor,
    input_shape: &[usize],
    axis: usize,
    keep_dim: bool,
    divisor: usize,
) -> Result<Tensor, TensorAutodiffError> {
    debug_assert!(divisor > 0);
    let mut result = zeros(input_shape)?;

    let mut upstream_axis = 0;
    let source_strides = input_shape
        .iter()
        .enumerate()
        .map(|(input_axis, _)| {
            if input_axis == axis {
                if keep_dim {
                    upstream_axis += 1;
                }
                0
            } else {
                let stride = upstream.strides()[upstream_axis];
                upstream_axis += 1;
                stride
            }
        })
        .collect::<Vec<_>>();
    debug_assert_eq!(upstream_axis, upstream.rank());
    let source_offsets = upstream
        .view()
        .projected_offsets(input_shape, &source_strides, result.len())
        .expect("a checked reduction VJP retains a valid source traversal plan");

    for (destination, source_offset) in result.as_mut_slice().iter_mut().zip(source_offsets) {
        *destination = upstream.as_slice()[source_offset] / divisor as f64;
    }
    Ok(result)
}

backward() is the ordinary call for a tracked rank-zero output; it supplies a scalar seed of one. A non-scalar output instead uses backward_with_seed with a finite tensor of exactly the same shape. Both calls build the node order, hold fresh adjoints for the duration of the pass, and read the saved context required by each local VJP. They commit parameter gradients but do not create a node-and- edge report.

When a caller needs to inspect that report, backward_with_trace or backward_with_seed_and_trace observes the same reverse calculation and returns TensorBackwardPass. For every edge, the trace records its upstream adjoint, local contribution, and the pass-local parent adjoint immediately before and after adding that contribution. Each node record also shows its pass-local adjoint and the validated prospective stored gradient for a parameter leaf. The trace does not recompute the derivatives.

backward() and backward_with_trace() always retain the graph. The seeded methods accept an explicit GraphRetention choice, independently of whether a trace is requested. To release a scalar graph, call a seeded method with the rank-zero seed one and GraphRetention::Release. Reverse traversal visits each node once but processes every ordered edge, including both multiply edges. No stored gradient changes unless all VJPs, pass accumulators, and prospective parameter sums validate.

Backward first verifies that the selected output still has graph context and is tracked. It then builds the reachable topology, rejecting any reachable operation whose graph context was released. Next it checks that the seed has exactly the output shape and contains only finite values. Only then does it compare the recorded and current parent revisions on every reachable operand edge. It checks edges in deterministic reverse-topological order and, within each operation, in operand order. This revision scan finishes before backward applies any VJP, asks the optional trace observer to record an edge, acquires a gradient write guard, or releases graph context. A mismatch returns StaleOperandValue, identifying the child, parent, operand slot, recorded revision, and current revision. The rejected call changes no stored gradient and does not release any operation context. The caller must run a new forward pass to build edges and saved context from the updated parameter values.

Accumulate fresh tensor adjoints transactionally, then retain or release saved operation context rust/crates/llm-from-scratch/src/autograd/tensor_core.rs#tensor-reverse-pass

    /// Reverses a rank-zero output with an implicit scalar seed of one.
    pub fn backward(&self) -> Result<(), TensorAutodiffError> {
        self.backward_scalar(NoTensorBackwardTrace)
    }

    /// Reverses a rank-zero output and records its node and edge evidence.
    pub fn backward_with_trace(&self) -> Result<TensorBackwardPass, TensorAutodiffError> {
        self.backward_scalar(RecordTensorBackwardTrace::default())
    }

    fn backward_scalar<Observer: TensorBackwardObserver>(
        &self,
        observer: Observer,
    ) -> Result<Observer::Output, TensorAutodiffError> {
        if self.is_released() {
            return Err(TensorAutodiffError::GraphReleased {
                operation: self.operation(),
            });
        }
        if !self.tracks_gradient() {
            return Err(TensorAutodiffError::UntrackedOutput {
                operation: self.operation(),
            });
        }
        if self.shape() != Vec::<usize>::new() {
            return Err(TensorAutodiffError::SeedShapeMismatch {
                expected: self.shape(),
                actual: Vec::new(),
            });
        }
        let seed = Tensor::from_vec(Vec::new(), vec![1.0])?;
        self.backward_with_observer(&seed.view(), GraphRetention::Retain, observer)
    }

    /// Runs a fresh exact-shape reverse pass without creating a trace record.
    ///
    /// Stored parameter gradients and graph edges remain bit-identical unless
    /// every VJP, pass accumulation, and prospective stored gradient is finite.
    pub fn backward_with_seed(
        &self,
        seed: &TensorView<'_>,
        retention: GraphRetention,
    ) -> Result<(), TensorAutodiffError> {
        self.backward_with_observer(seed, retention, NoTensorBackwardTrace)
    }

    /// Runs a fresh exact-shape reverse pass and records its trace.
    pub fn backward_with_seed_and_trace(
        &self,
        seed: &TensorView<'_>,
        retention: GraphRetention,
    ) -> Result<TensorBackwardPass, TensorAutodiffError> {
        self.backward_with_observer(seed, retention, RecordTensorBackwardTrace::default())
    }

    fn backward_with_observer<Observer: TensorBackwardObserver>(
        &self,
        seed: &TensorView<'_>,
        retention: GraphRetention,
        mut observer: Observer,
    ) -> Result<Observer::Output, TensorAutodiffError> {
        if self.is_released() {
            return Err(TensorAutodiffError::GraphReleased {
                operation: self.operation(),
            });
        }
        if !self.tracks_gradient() {
            return Err(TensorAutodiffError::UntrackedOutput {
                operation: self.operation(),
            });
        }
        let topology = self.topology()?;
        let expected = self.shape();
        if seed.shape() != expected {
            return Err(TensorAutodiffError::SeedShapeMismatch {
                expected,
                actual: seed.shape().to_vec(),
            });
        }
        let seed = seed.materialize()?;
        if let Some((index, value)) = first_nonfinite(&seed) {
            return Err(TensorAutodiffError::NonFiniteSeed { index, value });
        }

        let indices = topology
            .iter()
            .enumerate()
            .map(|(index, value)| (value.key(), index))
            .collect::<HashMap<_, _>>();
        for (child, value) in topology.iter().enumerate().rev() {
            let state = value.node.state.borrow();
            for (operand, edge) in state.parents.iter().enumerate() {
                let current_revision = edge.parent.value_revision();
                if edge.parent_value_revision != current_revision {
                    return Err(TensorAutodiffError::StaleOperandValue {
                        child,
                        parent: indices[&edge.parent.key()],
                        operand,
                        recorded_revision: edge.parent_value_revision,
                        current_revision,
                    });
                }
            }
        }
        let mut pass_adjoints = vec![None; topology.len()];
        pass_adjoints[topology.len() - 1] = Some(seed);

        for child in (0..topology.len()).rev() {
            let Some(upstream) = pass_adjoints[child].clone() else {
                continue;
            };
            let state = topology[child].node.state.borrow();
            for (operand, edge) in state.parents.iter().enumerate() {
                let parent = indices[&edge.parent.key()];
                let contribution = apply_vjp(&upstream, &edge.saved)?;
                if let Some((index, value)) = first_nonfinite(&contribution) {
                    return Err(TensorAutodiffError::NonFiniteVjp {
                        child,
                        parent,
                        operand,
                        index,
                        value,
                    });
                }

                let parent_tracked = edge.parent.tracks_gradient();
                if parent_tracked {
                    let previous = pass_adjoints[parent]
                        .clone()
                        .unwrap_or(zeros(edge.parent.shape().as_slice())?);
                    let next =
                        add_checked(&previous, &contribution, |index, previous, contribution| {
                            TensorAutodiffError::NonFinitePassAdjoint {
                                node: parent,
                                index,
                                previous,
                                contribution,
                            }
                        })?;
                    observer.observe_edge(
                        child,
                        parent,
                        operand,
                        &edge.saved,
                        &upstream,
                        &contribution,
                        true,
                        Some(&previous),
                        Some(&next),
                    );
                    pass_adjoints[parent] = Some(next);
                } else {
                    observer.observe_edge(
                        child,
                        parent,
                        operand,
                        &edge.saved,
                        &upstream,
                        &contribution,
                        false,
                        None,
                        None,
                    );
                }
            }
        }

        let mut prospective = vec![None; topology.len()];
        for (index, value) in topology.iter().enumerate() {
            let Some(stored) = value.gradient() else {
                continue;
            };
            let pass = pass_adjoints[index]
                .clone()
                .unwrap_or(zeros(value.shape().as_slice())?);
            prospective[index] = Some(add_checked(
                &stored,
                &pass,
                |element, stored, pass_adjoint| TensorAutodiffError::NonFiniteAccumulatedGradient {
                    node: index,
                    index: element,
                    stored,
                    pass_adjoint,
                },
            )?);
        }

        let observation = observer.finish(retention, &topology, &pass_adjoints, &prospective);

        let mut commits: Vec<(RefMut<'_, NodeState>, Tensor)> = Vec::new();
        for (value, gradient) in topology.iter().zip(prospective) {
            if let Some(gradient) = gradient {
                let state = value
                    .node
                    .state
                    .try_borrow_mut()
                    .map_err(|_| TensorAutodiffError::GradientBorrowed)?;
                commits.push((state, gradient));
            }
        }
        for (mut state, gradient) in commits {
            state.parameter_gradient = Some(gradient);
        }

        if retention == GraphRetention::Release {
            for value in &topology {
                if !value.operation().is_leaf() {
                    let mut state = value.node.state.borrow_mut();
                    state.parents.clear();
                    state.released = true;
                }
            }
        }

        Ok(observation)
    }

    /// Clears this parameter's accumulated gradient without changing its tape.
    pub fn zero_grad(&self) -> Result<(), TensorAutodiffError> {
        if !self.is_parameter() {
            return Err(TensorAutodiffError::NotAParameter {
                operation: self.operation(),
            });
        }
        let shape = self.shape();
        self.node
            .state
            .try_borrow_mut()
            .map_err(|_| TensorAutodiffError::GradientBorrowed)?
            .parameter_gradient = Some(zeros(&shape)?);
        Ok(())
    }

The worked example explicitly requests a trace for its first pass because it inspects the edge-by-edge evidence shown below. Its second pass needs only the committed gradients and therefore uses the ordinary call. The parameter primals do not change between these calls, so every captured revision still matches. The retained second pass recomputes fresh pass-local adjoints and contributions and doubles the two stored parameter gradients. Retention keeps saved context available, but it does not make that context valid across an in-place parameter update. zero_grad writes positive zero through saved parameter handles. A successful releasing pass recomputes fresh pass-local adjoints after zeroing, commits values that are numerically equal to the first-pass gradients, then drops the reachable operation nodes’ parent edges and saved context. Primal values, committed parameter gradients, and an explicit trace returned earlier remain readable, but another backward call or differentiable reuse of a released operation fails without mutation.

The short example sets parameter $p=[2,3]$ and adds $p\odot p$ to $10\operatorname{detach}(p)$ before summing. The two branches contribute $13$ and $50$ to the forward total $63$ , but the detached branch has no edge to $p$ , so the gradient of $p$ is [4,6]. The same example compares every supported VJP with Chapter 13 sampled central differences:

Check retention, zeroing, release, detached sum behavior, and every supported VJP rust/demos/ch15-tensor-autodiff-core/src/lib.rs#tensor-autodiff-lifecycle-gradcheck

/// Keeps one live branch and stops the equal-valued detached branch.
pub fn detach_sum_example() -> Result<DetachSumExample, TensorAutodiffError> {
    let p = TensorValue::parameter(tensor(&[2], &[2.0, 3.0]))?;
    let squared = p.mul(&p)?;
    let detached = p.detach();
    let ten = TensorValue::constant(tensor(&[], &[10.0]))?;
    let stopped = detached.mul(&ten)?;
    let elements = squared.add(&stopped)?;
    let output = elements.sum_axis(0, false)?;
    output.backward()?;

    Ok(DetachSumExample {
        value: output.value().as_slice()[0],
        p_gradient: p.gradient_snapshot().expect("p is a parameter"),
        detached_gradient: detached.gradient_snapshot(),
        detached_is_new_node: !p.is_same_node(&detached),
        detached_tracks_gradient: detached.tracks_gradient(),
    })
}

/// Checks add, multiply, reshape, transpose, broadcast, sum, and mean VJPs.
pub fn vjp_gradcheck_example() -> Result<VjpGradcheckExample, Box<dyn Error>> {
    let matrix_seed = tensor(&X_SHAPE, &[1.0, -2.0, 0.5, 3.0, -1.5, 2.0]);
    let bias = frozen_bias();
    let add = sampled_vjp_check(
        frozen_x(),
        &matrix_seed,
        |parameter| {
            let bias = TensorValue::constant(bias.clone())?;
            parameter.add(&bias)
        },
        |candidate| {
            let values = candidate
                .as_slice()
                .iter()
                .enumerate()
                .map(|(offset, value)| value + BIAS_VALUES[offset % X_SHAPE[1]])
                .collect::<Vec<_>>();
            weighted_sum(&values, matrix_seed.as_slice())
        },
    )?;

    let multiply = sampled_vjp_check(
        frozen_x(),
        &matrix_seed,
        |parameter| {
            let bias = TensorValue::constant(bias.clone())?;
            parameter.mul(&bias)
        },
        |candidate| {
            let values = candidate
                .as_slice()
                .iter()
                .enumerate()
                .map(|(offset, value)| value * BIAS_VALUES[offset % X_SHAPE[1]])
                .collect::<Vec<_>>();
            weighted_sum(&values, matrix_seed.as_slice())
        },
    )?;

    let reshape_seed = tensor(&RESHAPED_SHAPE, &[0.5, 1.0, -1.0, 2.0, 3.0, -0.25]);
    let reshape = sampled_vjp_check(
        frozen_x(),
        &reshape_seed,
        |parameter| parameter.reshape(&RESHAPED_SHAPE),
        |candidate| weighted_sum(candidate.as_slice(), reshape_seed.as_slice()),
    )?;

    let transpose_seed = tensor(&[3, 2], &[1.0, -2.0, 3.0, -4.0, 5.0, -6.0]);
    let transpose = sampled_vjp_check(
        frozen_x(),
        &transpose_seed,
        |parameter| parameter.transpose(0, 1),
        |candidate| {
            let mut result = 0.0;
            for row in 0..X_SHAPE[0] {
                for column in 0..X_SHAPE[1] {
                    let input = row * X_SHAPE[1] + column;
                    let output = column * X_SHAPE[0] + row;
                    result += candidate.as_slice()[input] * transpose_seed.as_slice()[output];
                }
            }
            result
        },
    )?;

    let broadcast = sampled_vjp_check(
        frozen_bias(),
        &matrix_seed,
        |parameter| parameter.broadcast_to(&X_SHAPE),
        |candidate| {
            let mut result = 0.0;
            for row in 0..X_SHAPE[0] {
                for column in 0..X_SHAPE[1] {
                    result += candidate.as_slice()[column]
                        * matrix_seed.as_slice()[row * X_SHAPE[1] + column];
                }
            }
            result
        },
    )?;

    let reduction_seed = tensor(&OUTPUT_SHAPE, &[2.0, -3.0]);
    let sum = sampled_vjp_check(
        frozen_x(),
        &reduction_seed,
        |parameter| parameter.sum_axis(1, false),
        |candidate| {
            (0..X_SHAPE[0])
                .map(|row| {
                    candidate.as_slice()[row * X_SHAPE[1]..(row + 1) * X_SHAPE[1]]
                        .iter()
                        .sum::<f64>()
                        * reduction_seed.as_slice()[row]
                })
                .sum()
        },
    )?;
    let mean = sampled_vjp_check(
        frozen_x(),
        &reduction_seed,
        |parameter| parameter.mean_axis(1, false),
        |candidate| {
            (0..X_SHAPE[0])
                .map(|row| {
                    candidate.as_slice()[row * X_SHAPE[1]..(row + 1) * X_SHAPE[1]]
                        .iter()
                        .sum::<f64>()
                        / X_SHAPE[1] as f64
                        * reduction_seed.as_slice()[row]
                })
                .sum()
        },
    )?;

    let passed = [
        &add, &multiply, &reshape, &transpose, &broadcast, &sum, &mean,
    ]
    .into_iter()
    .all(|check| check.passed);
    Ok(VjpGradcheckExample {
        add,
        multiply,
        reshape,
        transpose,
        broadcast,
        sum,
        mean,
        passed,
    })
}

/// Exercises typed seed, release, and transactional accumulation failures.
pub fn typed_error_example() -> Result<TypedErrorExample, TensorAutodiffError> {
    let shape_parameter = TensorValue::parameter(tensor(&[2], &[2.0, 3.0]))?;
    let shape_output = shape_parameter.mul(&shape_parameter)?;
    let shape_before = optional_tensor_bits(shape_parameter.gradient());
    let wrong_shape = tensor(&[1], &[1.0]);
    let seed_shape_error = shape_output
        .backward_with_seed(&wrong_shape.view(), GraphRetention::Release)
        .expect_err("a wrong-shape seed must fail");
    let seed_shape_unchanged = optional_tensor_bits(shape_parameter.gradient()) == shape_before;
    let valid_shape_seed = tensor(&[2], &[1.0, 1.0]);
    let seed_shape_graph_unchanged = !shape_output.is_released()
        && shape_output
            .backward_with_seed(&valid_shape_seed.view(), GraphRetention::Retain)
            .is_ok();

    let finite_parameter = TensorValue::parameter(tensor(&[2], &[2.0, 3.0]))?;
    let finite_output = finite_parameter.mul(&finite_parameter)?;
    let finite_before = optional_tensor_bits(finite_parameter.gradient());
    let nonfinite = tensor(&[2], &[1.0, f64::NAN]);
    let nonfinite_seed_error = finite_output
        .backward_with_seed(&nonfinite.view(), GraphRetention::Release)
        .expect_err("a non-finite seed must fail");
    let nonfinite_seed_unchanged =
        optional_tensor_bits(finite_parameter.gradient()) == finite_before;
    let nonfinite_graph_unchanged = !finite_output.is_released()
        && finite_output
            .backward_with_seed(&valid_shape_seed.view(), GraphRetention::Retain)
            .is_ok();

    let released_parameter = TensorValue::parameter(tensor(&[2], &[2.0, 3.0]))?;
    let released_square = released_parameter.mul(&released_parameter)?;
    let released_output = released_square.mean_axis(0, false)?;
    released_output.backward_with_seed(&tensor(&[], &[1.0]).view(), GraphRetention::Release)?;
    let released_before = optional_tensor_bits(released_parameter.gradient());
    let graph_released_error = released_output
        .backward()
        .expect_err("a released mean must reject another pass");
    let graph_released_unchanged =
        optional_tensor_bits(released_parameter.gradient()) == released_before;
    let graph_still_released = released_output.is_released();

    let accumulated = TensorValue::parameter(tensor(&[1], &[1.0]))?;
    let maximum = tensor(&[1], &[f64::MAX]);
    accumulated.backward_with_seed(&maximum.view(), GraphRetention::Retain)?;
    let accumulated_before = optional_tensor_bits(accumulated.gradient());
    let nonfinite_accumulation_error = accumulated
        .backward_with_seed(&maximum.view(), GraphRetention::Release)
        .expect_err("overflowing a stored parameter gradient must fail");
    let nonfinite_accumulation_unchanged =
        optional_tensor_bits(accumulated.gradient()) == accumulated_before;
    let cancellation = tensor(&[1], &[-f64::MAX]);
    let accumulation_graph_unchanged = !accumulated.is_released()
        && accumulated
            .backward_with_seed(&cancellation.view(), GraphRetention::Retain)
            .is_ok();

    let seed_shape = AtomicTensorError {
        error: seed_shape_error,
        gradients_unchanged: seed_shape_unchanged,
        graph_unchanged: seed_shape_graph_unchanged,
    };
    let nonfinite_seed = AtomicTensorError {
        error: nonfinite_seed_error,
        gradients_unchanged: nonfinite_seed_unchanged,
        graph_unchanged: nonfinite_graph_unchanged,
    };
    let graph_released = AtomicTensorError {
        error: graph_released_error,
        gradients_unchanged: graph_released_unchanged,
        graph_unchanged: graph_still_released,
    };
    let nonfinite_accumulation = AtomicTensorError {
        error: nonfinite_accumulation_error,
        gradients_unchanged: nonfinite_accumulation_unchanged,
        graph_unchanged: accumulation_graph_unchanged,
    };
    let all_unchanged = [
        &seed_shape,
        &nonfinite_seed,
        &graph_released,
        &nonfinite_accumulation,
    ]
    .into_iter()
    .all(|failure| failure.gradients_unchanged && failure.graph_unchanged);

    Ok(TypedErrorExample {
        seed_shape,
        nonfinite_seed,
        graph_released,
        nonfinite_accumulation,
        all_unchanged,
    })
}

Typed errors distinguish unsafe leaves and results, invalid operations, seed shape or value failures, non-finite reverse values, untracked outputs, released graphs, a changed parent primal reached through an old retained graph, and a stored gradient that is still read-borrowed. A failed pass leaves stored gradient bits and lifecycle state unchanged.

Reject unsafe tensor autodiff requests without partial gradient or graph mutation rust/crates/llm-from-scratch/src/autograd/tensor_core.rs#tensor-autodiff-errors

/// A deterministic rejection from tensor tape construction or reversal.
#[derive(Clone, Debug, PartialEq)]
pub enum TensorAutodiffError {
    Tensor(TensorError),
    View(TensorViewError),
    Operation(TensorOpError),
    Matmul(MatmulError),
    Probability(ProbabilityError),
    Model(ModelOpError),
    BroadcastTargetMismatch {
        input: Vec<usize>,
        requested: Vec<usize>,
        inferred: Vec<usize>,
    },
    NonFiniteLeaf {
        operation: TensorOperation,
        index: usize,
        value: f64,
    },
    NonFiniteForward {
        operation: TensorOperation,
        index: usize,
        value: f64,
    },
    UntrackedOutput {
        operation: TensorOperation,
    },
    GraphReleased {
        operation: TensorOperation,
    },
    ReleasedOperand {
        operation: TensorOperation,
        operand: usize,
    },
    SeedShapeMismatch {
        expected: Vec<usize>,
        actual: Vec<usize>,
    },
    NonFiniteSeed {
        index: usize,
        value: f64,
    },
    StaleOperandValue {
        child: usize,
        parent: usize,
        operand: usize,
        recorded_revision: u64,
        current_revision: u64,
    },
    NonFiniteVjp {
        child: usize,
        parent: usize,
        operand: usize,
        index: usize,
        value: f64,
    },
    NonFinitePassAdjoint {
        node: usize,
        index: usize,
        previous: f64,
        contribution: f64,
    },
    NonFiniteAccumulatedGradient {
        node: usize,
        index: usize,
        stored: f64,
        pass_adjoint: f64,
    },
    GradientBorrowed,
    NotAParameter {
        operation: TensorOperation,
    },
}

impl fmt::Display for TensorAutodiffError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Tensor(error) => error.fmt(formatter),
            Self::View(error) => error.fmt(formatter),
            Self::Operation(error) => error.fmt(formatter),
            Self::Matmul(error) => error.fmt(formatter),
            Self::Probability(error) => error.fmt(formatter),
            Self::Model(error) => error.fmt(formatter),
            Self::BroadcastTargetMismatch {
                input,
                requested,
                inferred,
            } => write!(
                formatter,
                "cannot broadcast shape {input:?} exactly to {requested:?}; broadcasting infers {inferred:?}"
            ),
            Self::NonFiniteLeaf {
                operation,
                index,
                value,
            } => write!(
                formatter,
                "{operation} tensor value at flat index {index} must be finite, got {value:?}"
            ),
            Self::NonFiniteForward {
                operation,
                index,
                value,
            } => write!(
                formatter,
                "{operation} produced non-finite value {value:?} at flat index {index}"
            ),
            Self::UntrackedOutput { operation } => write!(
                formatter,
                "cannot backpropagate from untracked {operation} output"
            ),
            Self::GraphReleased { operation } => {
                write!(
                    formatter,
                    "the {operation} operation tape has been released"
                )
            }
            Self::ReleasedOperand { operation, operand } => write!(
                formatter,
                "cannot build {operation}: operand {operand} reaches a released operation tape"
            ),
            Self::SeedShapeMismatch { expected, actual } => write!(
                formatter,
                "backward seed shape {actual:?} does not match output shape {expected:?}"
            ),
            Self::NonFiniteSeed { index, value } => write!(
                formatter,
                "backward seed at flat index {index} must be finite, got {value:?}"
            ),
            Self::StaleOperandValue {
                child,
                parent,
                operand,
                recorded_revision,
                current_revision,
            } => write!(
                formatter,
                "cannot backpropagate through operand {operand} from topology node {child} to {parent}: the forward pass recorded parent value revision {recorded_revision}, but its current revision is {current_revision}; run a new forward pass"
            ),
            Self::NonFiniteVjp {
                child,
                parent,
                operand,
                index,
                value,
            } => write!(
                formatter,
                "edge {operand} from topology node {child} to {parent} produced non-finite VJP value {value:?} at flat index {index}"
            ),
            Self::NonFinitePassAdjoint {
                node,
                index,
                previous,
                contribution,
            } => write!(
                formatter,
                "topology node {node} cannot accumulate pass-adjoint value {previous:?} plus {contribution:?} at flat index {index}"
            ),
            Self::NonFiniteAccumulatedGradient {
                node,
                index,
                stored,
                pass_adjoint,
            } => write!(
                formatter,
                "topology node {node} cannot accumulate stored gradient {stored:?} plus pass adjoint {pass_adjoint:?} at flat index {index}"
            ),
            Self::GradientBorrowed => formatter.write_str(
                "cannot mutate a parameter gradient while a read-only gradient borrow is active",
            ),
            Self::NotAParameter { operation } => {
                write!(
                    formatter,
                    "cannot clear a gradient on {operation}; only parameters store gradients"
                )
            }
        }
    }
}

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

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

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

impl From<TensorOpError> for TensorAutodiffError {
    fn from(error: TensorOpError) -> Self {
        Self::Operation(error)
    }
}

impl From<MatmulError> for TensorAutodiffError {
    fn from(error: MatmulError) -> Self {
        Self::Matmul(error)
    }
}

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

impl From<ModelOpError> for TensorAutodiffError {
    fn from(error: ModelOpError) -> Self {
        Self::Model(error)
    }
}

The learner command prints the worked graph, lifecycle, focused checks, and errors in a stable order:

Prepare the deterministic tensor autodiff evidence before printing rust/demos/ch15-tensor-autodiff-core/src/main.rs#learner-tensor-autodiff-output

    let example = frozen_tensor_example()?;
    let detached = detach_sum_example()?;
    let gradchecks = vjp_gradcheck_example()?;
    let errors = typed_error_example()?;

./course run cargo run --quiet --locked -p ch15-tensor-autodiff-core

parameter x: shape=[2, 3] values=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
parameter bias: shape=[3] values=[1.0, -1.0, 0.0]
reshape: shape=[3, 2] values=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
transpose: shape=[2, 3] values=[1.0, 3.0, 5.0, 2.0, 4.0, 6.0]
broadcast: shape=[2, 3] values=[1.0, -1.0, 0.0, 1.0, -1.0, 0.0]
add: shape=[2, 3] values=[2.0, 2.0, 5.0, 3.0, 3.0, 6.0]
multiply reused: shape=[2, 3] values=[4.0, 4.0, 25.0, 9.0, 9.0, 36.0]
mean axis=1 keep_dim=false: shape=[2] values=[11.0, 18.0]
non-scalar seed: shape=[2] values=[3.0, 6.0]
one backward: x_grad=[4.0, 12.0, 4.0, 12.0, 10.0, 24.0] bias_grad=[16.0, 16.0, 34.0]
repeated backward: x_grad=[8.0, 24.0, 8.0, 24.0, 20.0, 48.0] bias_grad=[32.0, 32.0, 68.0]
zero_grad: x_grad=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0] bias_grad=[0.0, 0.0, 0.0]
after zero and release: x_grad=[4.0, 12.0, 4.0, 12.0, 10.0, 24.0] bias_grad=[16.0, 16.0, 34.0]
released graph: operation=mean gradients unchanged=true
detach and sum: value=63.0 p_grad=[4.0, 6.0] detached_grad=none
gradcheck: add | multiply | reshape | transpose | broadcast | sum | mean; pass=true
typed errors: seed-shape | non-finite-seed | graph-released | non-finite-accumulated-gradient; gradients unchanged=true
chapter 16 handoff: add model-critical tensor VJPs

Follow shape restoration in three focused views

The first figure keeps the forward tensor graph in topological order. Its eight node records expose the two parameter leaves, three structural operations, two elementwise operations, and final reduction, together with every forward shape, value, operation class, and pass-local adjoint. All eight operand-use records remain attached to their child operations, including both multiply operands.

The second figure carries seed [3,6] through the edge-major reverse ledger. Eight rows preserve reverse order 0 through 7 across mean, multiply, add, broadcast, transpose, and reshape. Each row keeps the upstream adjoint, saved context, reduced axes, exact parent shape, local rule, and resulting parent contribution.

The third figure separates pass-local adjoints from stored parameter gradients. It retains the first committed gradients, the doubled second retained pass, positive-zero state, and releasing commit equal to the first pass, then keeps sum, detach, sampled checks, and all four typed rejections together. Across all three figures, the presentation reads recorded relationships; the Rust implementation computes every shape, VJP, accumulated gradient, numerical comparison, lifecycle transition, and rejection.

Build the tensor graph without losing an operand use

Inspect eight forward nodes in topological order, with every operation, shape, value, pass-local adjoint, and all eight ordered operand edges.

Output tensor: [2] [11.000000000000, 18.000000000000]
Reverse seed: [2] [3.000000000000, 6.000000000000]
Unique graph nodes: 8
Operand edges: 8

Build one shape-changing tensor graph

Each operation owns a contiguous primal tensor. Node identity removes duplicate topology visits but retains both ordered edges from multiply back to the reused add result.

x

Forward order
0

Operation
parameter leaf

Shape
[2, 3]

Values
[1.000000000000, 2.000000000000, 3.000000000000, 4.000000000000, 5.000000000000, 6.000000000000]

Pass-local adjoint
[4.000000000000, 12.000000000000, 4.000000000000, 12.000000000000, 10.000000000000, 24.000000000000]
r

Forward order
1

Operation
reshape

Shape
[3, 2]

Values
[1.000000000000, 2.000000000000, 3.000000000000, 4.000000000000, 5.000000000000, 6.000000000000]

Pass-local adjoint
[4.000000000000, 12.000000000000, 4.000000000000, 12.000000000000, 10.000000000000, 24.000000000000]
- Operand edge 0 x
t

Forward order
2

Operation
transpose

Shape
[2, 3]

Values
[1.000000000000, 3.000000000000, 5.000000000000, 2.000000000000, 4.000000000000, 6.000000000000]

Pass-local adjoint
[4.000000000000, 4.000000000000, 10.000000000000, 12.000000000000, 12.000000000000, 24.000000000000]
- Operand edge 0 r
bias

Forward order
3

Operation
parameter leaf

Shape
[3]

Values
[1.000000000000, -1.000000000000, 0.000000000000]

Pass-local adjoint
[16.000000000000, 16.000000000000, 34.000000000000]
bb

Forward order
4

Operation
explicit broadcast

Shape
[2, 3]

Values
[1.000000000000, -1.000000000000, 0.000000000000, 1.000000000000, -1.000000000000, 0.000000000000]

Pass-local adjoint
[4.000000000000, 4.000000000000, 10.000000000000, 12.000000000000, 12.000000000000, 24.000000000000]
- Operand edge 0 bias
z

Forward order
5

Operation
elementwise add

Shape
[2, 3]

Values
[2.000000000000, 2.000000000000, 5.000000000000, 3.000000000000, 3.000000000000, 6.000000000000]

Pass-local adjoint
[4.000000000000, 4.000000000000, 10.000000000000, 12.000000000000, 12.000000000000, 24.000000000000]
- Operand edge 0 t
- Operand edge 1 bb
q

Forward order
6

Operation
elementwise multiply

Shape
[2, 3]

Values
[4.000000000000, 4.000000000000, 25.000000000000, 9.000000000000, 9.000000000000, 36.000000000000]

Pass-local adjoint
[1.000000000000, 1.000000000000, 1.000000000000, 2.000000000000, 2.000000000000, 2.000000000000]
- Operand edge 0 z
- Operand edge 1 z
y

Forward order
7

Operation
axis mean

Shape
[2]

Values
[11.000000000000, 18.000000000000]

Pass-local adjoint
[3.000000000000, 6.000000000000]
- Operand edge 0 q

Reverse the non-scalar seed through all eight edges

Follow seed [3, 6] from reverse order 0 through 7, pairing every edge with its upstream adjoint, local VJP, reduced axes, saved context, exact parent shape, and parent contribution.

Pull the non-scalar seed through every edge

Mean divides each component of the non-scalar seed by the saved extent three and repeats each quotient across one input row. Multiply sends one contribution through each repeated operand before add copies the completed adjoint to both parents.

Local VJP

axis mean: reinsert axis, broadcast, divide by extent
elementwise multiply: upstream times the other operand on each ordered edge
elementwise add: pass upstream to each operand, then sum to its shape
explicit broadcast: sum missing-leading and expanded-singleton axes
transpose: swap the saved axes again
reshape: restore the saved input shape

Reverse order	Child result	Operand edge	Parent input	Upstream adjoint	Local VJP	Reduced axes	Saved context	Parent contribution
0	`y [2]`	0	`q [2, 3]`	`[3.000000000000, 6.000000000000]`	axis mean	`1`	`Axis=1; Keep dimension=no; Divisor=3`	`[1.000000000000, 1.000000000000, 1.000000000000, 2.000000000000, 2.000000000000, 2.000000000000]`
1	`q [2, 3]`	0	`z [2, 3]`	`[1.000000000000, 1.000000000000, 1.000000000000, 2.000000000000, 2.000000000000, 2.000000000000]`	elementwise multiply	`none`	`Other operand [2, 3]=[2.000000000000, 2.000000000000, 5.000000000000, 3.000000000000, 3.000000000000, 6.000000000000]; Input shape=2x3; Output shape=2x3`	`[2.000000000000, 2.000000000000, 5.000000000000, 6.000000000000, 6.000000000000, 12.000000000000]`
2	`q [2, 3]`	1	`z [2, 3]`	`[1.000000000000, 1.000000000000, 1.000000000000, 2.000000000000, 2.000000000000, 2.000000000000]`	elementwise multiply	`none`	`Other operand [2, 3]=[2.000000000000, 2.000000000000, 5.000000000000, 3.000000000000, 3.000000000000, 6.000000000000]; Input shape=2x3; Output shape=2x3`	`[2.000000000000, 2.000000000000, 5.000000000000, 6.000000000000, 6.000000000000, 12.000000000000]`
3	`z [2, 3]`	0	`t [2, 3]`	`[4.000000000000, 4.000000000000, 10.000000000000, 12.000000000000, 12.000000000000, 24.000000000000]`	elementwise add	`none`	`none`	`[4.000000000000, 4.000000000000, 10.000000000000, 12.000000000000, 12.000000000000, 24.000000000000]`
4	`z [2, 3]`	1	`bb [2, 3]`	`[4.000000000000, 4.000000000000, 10.000000000000, 12.000000000000, 12.000000000000, 24.000000000000]`	elementwise add	`none`	`none`	`[4.000000000000, 4.000000000000, 10.000000000000, 12.000000000000, 12.000000000000, 24.000000000000]`
5	`bb [2, 3]`	0	`bias [3]`	`[4.000000000000, 4.000000000000, 10.000000000000, 12.000000000000, 12.000000000000, 24.000000000000]`	explicit broadcast	`0`	`none`	`[16.000000000000, 16.000000000000, 34.000000000000]`
6	`t [2, 3]`	0	`r [3, 2]`	`[4.000000000000, 4.000000000000, 10.000000000000, 12.000000000000, 12.000000000000, 24.000000000000]`	transpose	`none`	`Axes=0,1`	`[4.000000000000, 12.000000000000, 4.000000000000, 12.000000000000, 10.000000000000, 24.000000000000]`
7	`r [3, 2]`	0	`x [2, 3]`	`[4.000000000000, 12.000000000000, 4.000000000000, 12.000000000000, 10.000000000000, 24.000000000000]`	reshape	`none`	`Input shape=2x3; Output shape=3x2`	`[4.000000000000, 12.000000000000, 4.000000000000, 12.000000000000, 10.000000000000, 24.000000000000]`

Separate stored gradients, graph lifecycle, checks, and rejections

Compare retained accumulation, zeroing, release, detach, sampled checks, and four rejected requests without mixing pass-local adjoints with stored parameter gradients.

Restore exact parameter shapes

Broadcast sums the reused rows, transpose swaps its saved axes, and reshape restores the original shape. Every contribution exactly matches its parent.

first retained pass committed

Node: x
Shape: [2, 3]
Stored gradient: [4.000000000000, 12.000000000000, 4.000000000000, 12.000000000000, 10.000000000000, 24.000000000000]

first retained pass committed

Node: bias
Shape: [3]
Stored gradient: [16.000000000000, 16.000000000000, 34.000000000000]

Retain, accumulate, zero, release

Retained calls recompute fresh pass-local adjoints and accumulate only parameter gradients. Successful release keeps values and committed gradients while discarding operation context.

second retained pass accumulated

x: [8.000000000000, 24.000000000000, 8.000000000000, 24.000000000000, 20.000000000000, 48.000000000000]
bias: [32.000000000000, 32.000000000000, 68.000000000000]

parameter gradients zeroed

x: [0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000, 0.000000000000]
bias: [0.000000000000, 0.000000000000, 0.000000000000]

releasing pass recomputed and committed one-pass gradients

x: [4.000000000000, 12.000000000000, 4.000000000000, 12.000000000000, 10.000000000000, 24.000000000000]
bias: [16.000000000000, 16.000000000000, 34.000000000000]

operation graph released

Operation: mean
Released: yes
Gradients unchanged: yes

Check sum, detach, and every VJP

A focused sum-and-detach graph isolates those rules. Chapter 13 central differences probe every supported VJP independently of the reverse tape.

detached branch gradient path cut

Operation: sum(p*p+detach(p)*ten)
Values: 63.000000000000
Stored gradient: [4.000000000000, 6.000000000000]

sampled gradient checks passed

Operation: add, multiply, reshape, transpose, broadcast, sum, mean
x: Sampled flat coordinates: 0, 1, 3, 5
bias: Sampled flat coordinates: 0, 1, 2
Status: sampled gradient checks passed

Reject unsafe requests without mutation

The implementation validates every VJP, pass accumulator, prospective parameter sum, and release transition before mutation, so a failed request changes no gradient or graph state.

seed shape does not match the output

Expected shape=2; Actual shape=1

seed contains a non-finite value

Flat index=1; Value=nan

operation graph was already released

Operation=mean

prospective parameter gradient is not finite

Node=0; Flat index=0

Predict before running Rust

Predict the shape and values of every node from x and bias through y.
Count unique nodes and operand edges. Why are both totals eight even though multiply uses the add result twice?
Reverse mean from seed [3,6]. Give the effective source strides, source-offset sequence, expanded values, and both multiply-edge contributions.
For the broadcast-bias branch, give the effective destination strides and destination-offset sequence before summing into bias; then swap transpose axes and restore x’s shape.
Predict both parameter gradients after two retained passes, zeroing, and one releasing pass. Why is the second retained pass valid here, and what would an in-place parameter update between the two passes require?
An input has shape [2,3]; summing axis 0 with keep_dim=false produces shape [3]. With seed [10,20,30], give the effective source strides, source-offset sequence, and values returned to the two input rows.
Distinguish the value() and gradient() read guards, the value_snapshot() and gradient_snapshot() copies, detach, zero_grad, Retain, and Release by the data or state each one exposes or changes.
Misconception check: does a broadcast create independent parameter copies whose gradients can keep the output shape or select one occurrence?

Check the eight tensor-autodiff predictions

Reshape gives shape [3,2] with flat values unchanged. Transpose gives [1,3,5,2,4,6]; broadcast gives [1,-1,0,1,-1,0]; add gives [2,2,5,3,3,6]; multiply gives [4,4,25,9,9,36]; mean gives [11,18].
The nodes are x, reshape, transpose, bias, broadcast, add, multiply, and mean. The edges are one into reshape, one into transpose, one into broadcast, two into add, two into multiply, and one into mean. Node deduplication never removes the repeated multiply edge.
The mean input—the multiply result—has shape [2,3], and mean’s upstream tensor has shape [2]. Effective source strides [1,0] produce source offsets [0,0,0,1,1,1]. Reading seed [3,6] through those offsets and dividing by three gives [1,1,1,2,2,2] at multiply. Each multiply edge contributes upstream times add: [2,2,5,6,6,12]. Adding the two edges gives [4,4,10,12,12,24].
bias[3] aligns with the trailing axis of the [2,3] incoming contribution. Effective destination strides [0,1] produce offsets [0,1,2,0,1,2], so row-major additions give $\bar{\mathrm{bias}}=[4+12,4+12,10+24]=[16,16,34]$ . Swapping axes gives shape [3,2] and values [4,12,4,12,10,24]; reshape restores shape [2,3] with the same flat order.
First-pass gradients are $\bar{x}=[4,12,4,12,10,24]$ and $\bar{\mathrm{bias}}=[16,16,34]$ . The parameter primals and their revisions are unchanged, so the second retained pass is valid and doubles both gradients. Zeroing writes positive zero. One releasing pass recomputes and commits gradients equal to the first-pass values and then removes operation context. An in-place parameter update would advance the node’s current revision while the old edge retained its recorded revision: backward on the old graph would reject before any VJP, gradient write, or release, and a new forward pass would be required.
The original sum input has shape [2,3]. The upstream shape is [3], so effective source strides [0,1] produce offsets [0,1,2,0,1,2]. Reading the seed through those offsets returns [10,20,30] to the first row and the same [10,20,30] to the second row. Because this is sum rather than mean, the divisor is one.
value() and gradient() temporarily borrow node-owned data without copying it; the gradient guard exists only when a parameter has a stored gradient. value_snapshot() and gradient_snapshot() clone independent tensor data but create no tape node. A read guard must end before the corresponding storage can be mutated; an overlapping stored-gradient mutation returns GradientBorrowed without a partial commit. detach uses a primal snapshot to create a new untracked leaf. zero_grad clears one parameter’s stored gradient; Retain keeps operation context after a successful commit; Release discards that context after a successful commit but preserves primals and stored parameter gradients.
No. Broadcast reuses source coordinates. Its VJP sums every reused output coordinate back to the original parameter shape. Keeping [2,3] or selecting one row would both be wrong for bias[3].

After making the predictions, run the learner command shown above and compare its forward values and gradients with your calculations.

Prepare model-critical tensor gradients

The cumulative project can now reverse shape-preserving elementwise operations and structural tensor transformations while returning every contribution to its parent’s exact shape, accumulating only parameter-leaf gradients, and releasing saved operation context safely. VJPs that map one logical shape onto another can reuse the checked projected-stride traversal. Chapter 16 adds model-critical VJPs for matrix multiplication, repeated embedding gathers, nonlinearities, log-softmax, and indexed mean token loss; Chapter 15 alone still cannot train the decoder.

Every later decoder block will reuse this chapter’s graph contract: visit each operation once, retain every operand edge, apply one local VJP, add branches, commit parameter gradients transactionally, and release saved context only after success. Operations that undo broadcasting or reduction also reuse the structural projected-stride machinery. In Chapter 16, matrix multiplication uses broadcast reversal for parent contributions, while indexed token loss uses the checked offset cursor to visit group bases. Other operation-specific VJPs derive their own traversal plans from their own validated shapes.