Обратный проход по тензорному графу с локальными VJP

/// 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,
    })
}

#[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,
        )
    }
}

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)])
    }
}

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)
}

    /// 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(())
    }

/// 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,
    })
}

/// 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)
    }
}