zyx 0.17.0

// Copyright (C) 2025 zk4x
// SPDX-License-Identifier: LGPL-3.0-only WITH Classpath-exception-2.0
use crate::{
    Map, Set,
    dtype::Constant,
    graph::{Graph, GraphId, Node, OpId},
    kernel::{BOp, UOp},
    runtime::{Runtime, TensorData},
    shape::{Dim, UAxis},
    tensor::TensorId,
};
use std::collections::BTreeSet;

impl Runtime {
    pub(crate) fn gradient(&mut self, target: TensorId, sources: Set<TensorId>, graph_id: GraphId) -> Map<TensorId, TensorId> {
        let target_class = match self.tensors[target] {
            TensorData::Graph { class_id, .. }
            | TensorData::GraphLeaf { class_id, .. }
            | TensorData::Promoted { class_id, .. } => class_id,
            TensorData::Eager { .. } | TensorData::Leaf { .. } | TensorData::PendingLeaf { .. } | TensorData::Symbolic { .. } => {
                panic!("gradient on non-graph tensor")
            }
        };
        let source_classes: Set<OpId> = sources
            .iter()
            .map(|tid| match self.tensors[*tid] {
                TensorData::Graph { class_id, .. }
                | TensorData::GraphLeaf { class_id, .. }
                | TensorData::Promoted { class_id, .. } => class_id,
                TensorData::Eager { .. }
                | TensorData::Leaf { .. }
                | TensorData::PendingLeaf { .. }
                | TensorData::Symbolic { .. } => {
                    panic!("one of the sources is non-graph tensor: {tid}")
                }
            })
            .collect();

        let output_set: BTreeSet<OpId> = [target_class].into();
        let topo = self.graphs[graph_id].build_topo(&output_set, &source_classes);

        let mut grads: Map<OpId, OpId> = Map::default();

        // Seed gradient: ones expanded to the target's shape. Never a bare
        // const — a const class has no producer path, so it cannot be realized
        // as a tape output on its own. (Rank-0 targets keep the bare const;
        // the kernelizer stores consts like any other value.)
        let target_dtype = self.graphs[graph_id].dtype(target_class);
        let one_cid = self.push_const(graph_id, Constant::new(1u8).cast(target_dtype));
        let target_dims = self.graphs[graph_id].shape(target_class);
        let ones = if target_dims.is_empty() {
            one_cid
        } else {
            let shape_class = self.shape_class(graph_id, target_dims);
            self.push_node(graph_id, Node::Expand { x: one_cid, shape: shape_class }).1
        };
        grads.insert(target_class, ones);

        for &cid in &topo {
            let Some(&grad) = grads.get(&cid) else {
                continue;
            };

            let nid = cid;
            if matches!(&self.graphs[graph_id].nodes[nid].node, Node::Leaf { .. } | Node::Const { .. } | Node::Kernel { .. }) {
                continue;
            }

            match self.graphs[graph_id].nodes[nid].node {
                Node::Unary { x, uop } => match uop {
                    UOp::Neg => {
                        let g = self.push_node(graph_id, Node::Unary { x: grad, uop: UOp::Neg }).1;
                        accum_grad(self, graph_id, &mut grads, x, g);
                    }
                    UOp::Reciprocal => {
                        let z_sq = self.push_binary_node(graph_id, cid, cid, BOp::Mul);
                        let neg_z_sq = self.push_node(graph_id, Node::Unary { x: z_sq, uop: UOp::Neg }).1;
                        let g = self.push_binary_node(graph_id, grad, neg_z_sq, BOp::Mul);
                        accum_grad(self, graph_id, &mut grads, x, g);
                    }
                    UOp::Exp2 => {
                        let ln2_cid = self.push_const(graph_id, Constant::new(std::f64::consts::LN_2));
                        let z_ln2 = self.push_binary_node(graph_id, cid, ln2_cid, BOp::Mul);
                        let g = self.push_binary_node(graph_id, grad, z_ln2, BOp::Mul);
                        accum_grad(self, graph_id, &mut grads, x, g);
                    }
                    UOp::Log2 => {
                        let ln2_cid = self.push_const(graph_id, Constant::new(std::f64::consts::LN_2));
                        let x_ln2 = self.push_binary_node(graph_id, x, ln2_cid, BOp::Mul);
                        let g = self.push_binary_node(graph_id, grad, x_ln2, BOp::Div);
                        accum_grad(self, graph_id, &mut grads, x, g);
                    }
                    UOp::Sqrt => {
                        let two_cid = self.push_const(graph_id, Constant::new(2));
                        let z2 = self.push_binary_node(graph_id, cid, two_cid, BOp::Mul);
                        let g = self.push_binary_node(graph_id, grad, z2, BOp::Div);
                        accum_grad(self, graph_id, &mut grads, x, g);
                    }
                    UOp::Rsqrt => {
                        // dz/dx = -z^3/2 where z = rsqrt(x).
                        let z2 = self.push_binary_node(graph_id, cid, cid, BOp::Mul);
                        let z3 = self.push_binary_node(graph_id, z2, cid, BOp::Mul);
                        let two_cid = self.push_const(graph_id, Constant::new(2));
                        let z3_2 = self.push_binary_node(graph_id, z3, two_cid, BOp::Div);
                        let neg = self.push_node(graph_id, Node::Unary { x: z3_2, uop: UOp::Neg }).1;
                        let g = self.push_binary_node(graph_id, grad, neg, BOp::Mul);
                        accum_grad(self, graph_id, &mut grads, x, g);
                    }
                    UOp::Sin => {
                        let cos_x = self.push_node(graph_id, Node::Unary { x, uop: UOp::Cos }).1;
                        let g = self.push_binary_node(graph_id, grad, cos_x, BOp::Mul);
                        accum_grad(self, graph_id, &mut grads, x, g);
                    }
                    UOp::Cos => {
                        let sin_x = self.push_node(graph_id, Node::Unary { x, uop: UOp::Sin }).1;
                        let neg_sin = self.push_node(graph_id, Node::Unary { x: sin_x, uop: UOp::Neg }).1;
                        let g = self.push_binary_node(graph_id, grad, neg_sin, BOp::Mul);
                        accum_grad(self, graph_id, &mut grads, x, g);
                    }
                    UOp::Exp => {
                        let exp_x = self.push_node(graph_id, Node::Unary { x, uop: UOp::Exp }).1;
                        let g = self.push_binary_node(graph_id, grad, exp_x, BOp::Mul);
                        accum_grad(self, graph_id, &mut grads, x, g);
                    }
                    UOp::Abs => {
                        let zero = self.push_const(graph_id, Constant::new(0u8));
                        let one = self.push_const(graph_id, Constant::new(1u8));
                        let neg_one = self.push_const(graph_id, Constant::new(-1i8));
                        let is_pos = self.push_binary_node(graph_id, x, zero, BOp::Cmpgt);
                        let is_neg = self.push_binary_node(graph_id, x, zero, BOp::Cmplt);
                        let sign_pos = self.push_binary_node(graph_id, is_pos, one, BOp::Mul);
                        let sign_neg = self.push_binary_node(graph_id, is_neg, neg_one, BOp::Mul);
                        let sign = self.push_binary_node(graph_id, sign_pos, sign_neg, BOp::Add);
                        let g = self.push_binary_node(graph_id, grad, sign, BOp::Mul);
                        accum_grad(self, graph_id, &mut grads, x, g);
                    }
                    UOp::Floor | UOp::Trunc | UOp::BitNot | UOp::Not => {}
                },
                Node::Binary { x, y, bop } => match bop {
                    BOp::Add => {
                        accum_grad(self, graph_id, &mut grads, x, grad);
                        accum_grad(self, graph_id, &mut grads, y, grad);
                    }
                    BOp::Sub => {
                        accum_grad(self, graph_id, &mut grads, x, grad);
                        let neg_grad = self.push_node(graph_id, Node::Unary { x: grad, uop: UOp::Neg }).1;
                        accum_grad(self, graph_id, &mut grads, y, neg_grad);
                    }
                    BOp::Mul => {
                        let gx = self.push_binary_node(graph_id, grad, y, BOp::Mul);
                        accum_grad(self, graph_id, &mut grads, x, gx);
                        let gy = self.push_binary_node(graph_id, grad, x, BOp::Mul);
                        accum_grad(self, graph_id, &mut grads, y, gy);
                    }
                    BOp::Div => {
                        let gx = self.push_binary_node(graph_id, grad, y, BOp::Div);
                        accum_grad(self, graph_id, &mut grads, x, gx);
                        let neg_grad = self.push_node(graph_id, Node::Unary { x: grad, uop: UOp::Neg }).1;
                        let x_mul = self.push_binary_node(graph_id, neg_grad, x, BOp::Mul);
                        let y_sq = self.push_binary_node(graph_id, y, y, BOp::Mul);
                        let gy = self.push_binary_node(graph_id, x_mul, y_sq, BOp::Div);
                        accum_grad(self, graph_id, &mut grads, y, gy);
                    }
                    BOp::Pow => {
                        let one = self.push_const(graph_id, Constant::new(1u8));
                        let y_1 = self.push_binary_node(graph_id, y, one, BOp::Sub);
                        let x_pow_ym1 = self.push_binary_node(graph_id, x, y_1, BOp::Pow);
                        let y_mul = self.push_binary_node(graph_id, y, x_pow_ym1, BOp::Mul);
                        let gx = self.push_binary_node(graph_id, grad, y_mul, BOp::Mul);
                        accum_grad(self, graph_id, &mut grads, x, gx);
                        let log2_x = self.push_node(graph_id, Node::Unary { x, uop: UOp::Log2 }).1;
                        let ln2_cid = self.push_const(graph_id, Constant::new(std::f64::consts::LN_2));
                        let z_log2 = self.push_binary_node(graph_id, cid, log2_x, BOp::Mul);
                        let z_lnx = self.push_binary_node(graph_id, z_log2, ln2_cid, BOp::Mul);
                        let gy = self.push_binary_node(graph_id, grad, z_lnx, BOp::Mul);
                        accum_grad(self, graph_id, &mut grads, y, gy);
                    }
                    BOp::Mod => {
                        accum_grad(self, graph_id, &mut grads, x, grad);
                        let x_div_y = self.push_binary_node(graph_id, x, y, BOp::Div);
                        let floored = self.push_node(graph_id, Node::Unary { x: x_div_y, uop: UOp::Floor }).1;
                        let neg_floor = self.push_node(graph_id, Node::Unary { x: floored, uop: UOp::Neg }).1;
                        let gy = self.push_binary_node(graph_id, neg_floor, grad, BOp::Mul);
                        accum_grad(self, graph_id, &mut grads, y, gy);
                    }
                    BOp::Max => {
                        let x_gt_y = self.push_binary_node(graph_id, x, y, BOp::Cmpgt);
                        let x_lt_y = self.push_binary_node(graph_id, x, y, BOp::Cmplt);
                        let gx = self.push_binary_node(graph_id, grad, x_gt_y, BOp::Mul);
                        accum_grad(self, graph_id, &mut grads, x, gx);
                        let gy = self.push_binary_node(graph_id, grad, x_lt_y, BOp::Mul);
                        accum_grad(self, graph_id, &mut grads, y, gy);
                    }
                    BOp::Cmplt
                    | BOp::Cmpgt
                    | BOp::Cmpge
                    | BOp::Eq
                    | BOp::NotEq
                    | BOp::Or
                    | BOp::And
                    | BOp::BitXor
                    | BOp::BitOr
                    | BOp::BitAnd
                    | BOp::BitShiftLeft
                    | BOp::BitShiftRight => {}
                },
                Node::Cast { x, .. } => {
                    let g = self.push_node(graph_id, Node::Cast { x: grad, dtype: self.graphs[graph_id].dtype(x) }).1;
                    accum_grad(self, graph_id, &mut grads, x, g);
                }
                Node::Bitcast { x, .. } => {
                    let g = self.push_node(graph_id, Node::Bitcast { x: grad, dtype: self.graphs[graph_id].dtype(x) }).1;
                    accum_grad(self, graph_id, &mut grads, x, g);
                }
                Node::Reshape { x, .. } => {
                    let in_dims = self.graphs[graph_id].shape(x);
                    let x_shape = self.shape_class(graph_id, in_dims.clone());
                    let in_conc: Vec<Dim> = in_dims
                        .iter()
                        .map(|&d| self.graphs[graph_id].resolve_const(d).and_then(Constant::as_dim).unwrap_or(-1))
                        .collect();
                    let xc_conc: Vec<Dim> = self.graphs[graph_id]
                        .shape(x_shape)
                        .iter()
                        .map(|&d| self.graphs[graph_id].resolve_const(d).and_then(Constant::as_dim).unwrap_or(-1))
                        .collect();
                    if in_conc.iter().any(|&v| v != 0) && xc_conc.iter().any(|&v| v != 0) && in_conc != xc_conc {
                        eprintln!("RESGRAD in={:?} x_shape={:?}", in_conc, xc_conc);
                    }
                    let g = self.push_node(graph_id, Node::Reshape { x: grad, shape: x_shape }).1;
                    accum_grad(self, graph_id, &mut grads, x, g);
                }
                Node::Expand { x, .. } => {
                    let out_dims = self.graphs[graph_id].shape(cid);
                    let in_dims = self.graphs[graph_id].shape(x);
                    // Right-align the input against the expanded output per broadcast
                    // semantics. The input is broadcast to the output by (a) leading
                    // `pad` dims that the input did not have at all (implicitly size 1)
                    // and (b) trailing-aligned dims where the input is a singleton (1)
                    // broadcast to a larger output extent. The gradient of a broadcast
                    // must be summed over *all* of these axes to drop back to the
                    // input shape.
                    //
                    // Symbolic broadcast decision (tinygrad `broadcast_axes`
                    // semantics): an axis needs summing iff the input dim is
                    // provably a singleton (1); unknown symbolic dims default
                    // to NOT broadcast.
                    let pad = out_dims.len() - in_dims.len();
                    let mut sum_axes: Vec<UAxis> = (0..pad).map(|i| i as UAxis).collect();
                    for (i, &xd) in in_dims.iter().enumerate() {
                        if self.graph_const_dim(graph_id, xd) == Some(1) {
                            sum_axes.push((pad + i) as UAxis);
                        }
                    }
                    if sum_axes.is_empty() {
                        accum_grad(self, graph_id, &mut grads, x, grad);
                    } else {
                        let reduced_dims: Vec<OpId> = out_dims
                            .iter()
                            .enumerate()
                            .filter(|(i, _)| !sum_axes.contains(&(*i as UAxis)))
                            .map(|(_, &d)| d)
                            .collect();
                        let reduced =
                            self.push_node(graph_id, Node::Reduce { x: grad, rop: BOp::Add, axes: sum_axes.into_boxed_slice() });
                        // The graph reduce drops the reduced dims; restore the
                        // original shape (keepdim) with an explicit reshape.
                        let reduced = if reduced_dims == in_dims {
                            reduced.1
                        } else {
                            let xs = self.shape_class(graph_id, in_dims);
                            self.push_node(graph_id, Node::Reshape { x: reduced.1, shape: xs }).1
                        };
                        accum_grad(self, graph_id, &mut grads, x, reduced);
                    }
                }
                Node::Permute { x, ref axes } => {
                    let mut inv_axes: Vec<UAxis> = vec![0; axes.len()];
                    for (i, &a) in axes.iter().enumerate() {
                        inv_axes[a] = i as UAxis;
                    }
                    let g = self.push_node(graph_id, Node::Permute { x: grad, axes: inv_axes.into_boxed_slice() }).1;
                    accum_grad(self, graph_id, &mut grads, x, g);
                }
                Node::Pad { x, axis, lp, .. } => {
                    // Pad backward: narrow the gradient back to the original extent.
                    let orig_len = self.graphs[graph_id].shape(x)[axis as usize];
                    let g = self.push_node(graph_id, Node::Narrow { x: grad, axis, start: lp, len: orig_len }).1;
                    accum_grad(self, graph_id, &mut grads, x, g);
                }
                Node::Narrow { x, axis, start, .. } => {
                    // Narrow backward: pad the gradient with zeros back to the
                    // original extent.
                    let orig_len = self.graphs[graph_id].shape(x)[axis as usize];
                    let g = self.push_node(graph_id, Node::Pad { x: grad, axis, lp: start, len: orig_len }).1;
                    accum_grad(self, graph_id, &mut grads, x, g);
                }
                Node::Flip { x, ref axes } => {
                    // Flip is its own inverse: the gradient back-propagates by
                    // flipping along the same axes.
                    let g = self.push_node(graph_id, Node::Flip { x: grad, axes: axes.clone() }).1;
                    accum_grad(self, graph_id, &mut grads, x, g);
                }
                Node::Reduce { x, rop: bop, ref axes } => {
                    let axes = axes.clone();
                    match bop {
                        BOp::Add => {
                            // Reshape the gradient to x's dims with 1 at each
                            // reduced axis, then broadcast back to x's shape.
                            let x_dims = self.graphs[graph_id].shape(x);
                            // Shape dims are lengths: always integer-typed,
                            // never the tensor's data dtype.
                            let one_dim = self.push_const(graph_id, Constant::new(1i64));
                            let kept: Vec<OpId> = x_dims
                                .iter()
                                .enumerate()
                                .map(|(i, &d)| if axes.contains(&(i as UAxis)) { one_dim } else { d })
                                .collect();
                            let kept_shape = self.shape_class(graph_id, kept);
                            let grad_r = self.push_node(graph_id, Node::Reshape { x: grad, shape: kept_shape }).1;
                            let x_shape = self.shape_class(graph_id, x_dims);
                            let g = self.push_node(graph_id, Node::Expand { x: grad_r, shape: x_shape }).1;
                            accum_grad(self, graph_id, &mut grads, x, g);
                        }
                        BOp::Max => {
                            // Mask of positions that attained the max (1 - (x < z)),
                            // multiplied by the broadcast gradient.
                            let x_dims = self.graphs[graph_id].shape(x);
                            let dtype = self.graphs[graph_id].dtype(x);
                            let one = self.push_const(graph_id, Constant::new(1u8).cast(dtype));
                            // Shape dims are lengths: always integer-typed,
                            // never the tensor's data dtype.
                            let one_dim = self.push_const(graph_id, Constant::new(1i64));
                            let kept: Vec<OpId> = x_dims
                                .iter()
                                .enumerate()
                                .map(|(i, &d)| if axes.contains(&(i as UAxis)) { one_dim } else { d })
                                .collect();
                            let kept_shape = self.shape_class(graph_id, kept);
                            let x_shape = self.shape_class(graph_id, x_dims);
                            let z_reshaped = self.push_node(graph_id, Node::Reshape { x: cid, shape: kept_shape }).1;
                            let z_broadcasted = self.push_node(graph_id, Node::Expand { x: z_reshaped, shape: x_shape }).1;
                            let cmp = self.push_binary_node(graph_id, x, z_broadcasted, BOp::Cmplt);
                            let cmp_f = self.push_node(graph_id, Node::Cast { x: cmp, dtype }).1;
                            let one_e = self.push_node(graph_id, Node::Expand { x: one, shape: x_shape }).1;
                            let mask = self.push_binary_node(graph_id, one_e, cmp_f, BOp::Sub);
                            let grad_r = self.push_node(graph_id, Node::Reshape { x: grad, shape: kept_shape }).1;
                            let grad_e = self.push_node(graph_id, Node::Expand { x: grad_r, shape: x_shape }).1;
                            let grad_x = self.push_binary_node(graph_id, mask, grad_e, BOp::Mul);
                            accum_grad(self, graph_id, &mut grads, x, grad_x);
                        }
                        BOp::Mul => {
                            // d(prod x)/dx_i is the product of every *other* element.
                            // The short form grad * z / x breaks as soon as x holds a
                            // zero, so build it from the product of the nonzero
                            // elements (p) and the number of zeros (nz) along the
                            // reduced axes:
                            //   nz == 0 -> p / x_i
                            //   nz == 1 -> p at the zero itself, 0 everywhere else
                            //   nz > 1  -> 0
                            let x_dims = self.graphs[graph_id].shape(x);
                            let dtype = self.graphs[graph_id].dtype(x);
                            let zero = self.push_const(graph_id, Constant::new(0u8).cast(dtype));
                            let one = self.push_const(graph_id, Constant::new(1u8).cast(dtype));
                            // Shape dims are lengths: always integer-typed,
                            // never the tensor's data dtype.
                            let one_dim = self.push_const(graph_id, Constant::new(1i64));
                            let kept: Vec<OpId> = x_dims
                                .iter()
                                .enumerate()
                                .map(|(i, &d)| if axes.contains(&(i as UAxis)) { one_dim } else { d })
                                .collect();
                            let kept_shape = self.shape_class(graph_id, kept);
                            let x_shape = self.shape_class(graph_id, x_dims);
                            // 1 at every zero of x, 0 elsewhere.
                            let is_zero_b = self.push_binary_node(graph_id, x, zero, BOp::Eq);
                            let is_zero = self.push_node(graph_id, Node::Cast { x: is_zero_b, dtype }).1;
                            // Zeros become ones, so the product below keeps only the
                            // nonzero factors and the division never divides by zero.
                            let safe_x = self.push_binary_node(graph_id, x, is_zero, BOp::Add);
                            let p = self.push_node(graph_id, Node::Reduce { x: safe_x, rop: BOp::Mul, axes: axes.clone() }).1;
                            let nz = self.push_node(graph_id, Node::Reduce { x: is_zero, rop: BOp::Add, axes: axes.clone() }).1;
                            let p_r = self.push_node(graph_id, Node::Reshape { x: p, shape: kept_shape }).1;
                            let p_e = self.push_node(graph_id, Node::Expand { x: p_r, shape: x_shape }).1;
                            let nz_r = self.push_node(graph_id, Node::Reshape { x: nz, shape: kept_shape }).1;
                            let nz_e = self.push_node(graph_id, Node::Expand { x: nz_r, shape: x_shape }).1;
                            // No zero along the axis: every element divides p.
                            let no_zero_b = self.push_binary_node(graph_id, nz_e, zero, BOp::Eq);
                            let no_zero = self.push_node(graph_id, Node::Cast { x: no_zero_b, dtype }).1;
                            let quot = self.push_binary_node(graph_id, p_e, safe_x, BOp::Div);
                            let dense = self.push_binary_node(graph_id, quot, no_zero, BOp::Mul);
                            // Exactly one zero along the axis: p goes to that element.
                            let one_zero_b = self.push_binary_node(graph_id, nz_e, one, BOp::Eq);
                            let one_zero = self.push_node(graph_id, Node::Cast { x: one_zero_b, dtype }).1;
                            let at_zero = self.push_binary_node(graph_id, one_zero, is_zero, BOp::Mul);
                            let sparse = self.push_binary_node(graph_id, p_e, at_zero, BOp::Mul);
                            let partial = self.push_binary_node(graph_id, dense, sparse, BOp::Add);
                            let grad_r = self.push_node(graph_id, Node::Reshape { x: grad, shape: kept_shape }).1;
                            let grad_e = self.push_node(graph_id, Node::Expand { x: grad_r, shape: x_shape }).1;
                            let grad_x = self.push_binary_node(graph_id, partial, grad_e, BOp::Mul);
                            accum_grad(self, graph_id, &mut grads, x, grad_x);
                        }
                        ref bop => todo!("gradient for reduce {bop:?} is not yet supported"),
                    }
                }
                Node::ToDevice { x, .. } => {
                    accum_grad(self, graph_id, &mut grads, x, grad);
                }
                Node::Contiguous { x } => {
                    accum_grad(self, graph_id, &mut grads, x, grad);
                }
                Node::Assign { dst: _, src, .. } => {
                    accum_grad(self, graph_id, &mut grads, src, grad);
                }
                Node::After { x, .. } => {
                    accum_grad(self, graph_id, &mut grads, x, grad);
                }
                Node::Stack { .. } => todo!("stack backward"),
                Node::Index { vec, .. } => {
                    // Selecting one output of a multi-output kernel passes the
                    // gradient through to the Stack class.
                    accum_grad(self, graph_id, &mut grads, vec, grad);
                }
                Node::Leaf { .. } | Node::Const { .. } => {}
                Node::Kernel { .. } => todo!("backward through custom kernel"),
                Node::Custom { .. } => todo!("backward through custom kernel"),
            }
        }

        grads.retain(|k, _| source_classes.contains(k));

        let mut res = Map::default();
        for tid in sources {
            // The gradient result shares the source's shape expression.
            let (shape_id, dtype) = match self.tensors[tid] {
                TensorData::Graph { shape_id, dtype, .. }
                | TensorData::GraphLeaf { shape_id, dtype, .. }
                | TensorData::Promoted { shape_id, dtype, .. } => (shape_id, dtype),
                TensorData::Eager { .. }
                | TensorData::Leaf { .. }
                | TensorData::PendingLeaf { .. }
                | TensorData::Symbolic { .. } => {
                    panic!("gradient source {tid} is not a graph tensor: {:?}", self.tensors[tid])
                }
            };
            // Shape expressions live in the append-only expr slab: shared by
            // reference, no retain needed.
            let _ = shape_id;
            let grad_tid = match grads.get(&match self.tensors[tid] {
                TensorData::Graph { class_id, .. }
                | TensorData::GraphLeaf { class_id, .. }
                | TensorData::Promoted { class_id, .. } => class_id,
                TensorData::Eager { .. }
                | TensorData::Leaf { .. }
                | TensorData::PendingLeaf { .. }
                | TensorData::Symbolic { .. } => {
                    unreachable!("{:?}", self.tensors[tid])
                }
            }) {
                Some(&gcid) => gcid,
                None => {
                    let shape: Vec<Dim> = self.resolve_shape(tid);
                    let dtype = self.dtype(tid);
                    let zero_cid = self.push_const(graph_id, Constant::new(0u8).cast(dtype));
                    let ops: Box<[OpId]> = shape.iter().map(|&d| self.push_const(graph_id, Constant::idx(d))).collect();
                    let shape_cid = if ops.len() == 1 {
                        ops[0]
                    } else {
                        self.push_node(graph_id, Node::Stack { ops }).1
                    };
                    self.push_node(graph_id, Node::Expand { x: zero_cid, shape: shape_cid }).1
                }
            };
            self.graphs[graph_id].ref_count += 1;
            let grad_tid = self.tensors.push(TensorData::Graph { class_id: grad_tid, graph_id, shape_id, dtype, rc: 1 });
            res.insert(tid, grad_tid);
        }
        res
    }
}

impl Graph {
    pub fn build_topo(&self, outputs: &BTreeSet<OpId>, sources: &Set<OpId>) -> Vec<OpId> {
        let mut stack: Vec<OpId> = outputs.iter().copied().collect();
        let mut rcs: Map<OpId, u32> = Map::default();
        while let Some(cid) = stack.pop() {
            rcs.entry(cid).and_modify(|rc| *rc += 1).or_insert_with(|| {
                let nid = cid;
                let node = &self.nodes[nid].node;
                if !matches!(
                    node,
                    Node::Binary {
                        bop: BOp::Cmpgt
                            | BOp::Cmplt
                            | BOp::Eq
                            | BOp::NotEq
                            | BOp::Or
                            | BOp::And
                            | BOp::BitAnd
                            | BOp::BitOr
                            | BOp::BitXor
                            | BOp::BitShiftLeft
                            | BOp::BitShiftRight,
                        ..
                    }
                ) {
                    for p in node.class_params() {
                        if !stack.contains(&p) {
                            stack.push(p);
                        }
                    }
                }
                1
            });
        }

        let mut order = Vec::new();
        let mut internal_rcs: Map<OpId, u32> = Map::default();
        let mut stack: Vec<OpId> = outputs.iter().copied().collect();
        while let Some(cid) = stack.pop() {
            if let Some(&rc) = rcs.get(&cid)
                && rc == *internal_rcs.entry(cid).and_modify(|c| *c += 1).or_insert(1)
            {
                order.push(cid);
                for p in self.nodes[cid].node.class_params() {
                    if !stack.contains(&p) {
                        stack.push(p);
                    }
                }
            }
        }

        let mut topo = Vec::new();
        let mut req_grad = sources.clone();
        let mut visited: Set<OpId> = Set::default();
        for cid in order.into_iter().rev() {
            for p in self.nodes[cid].node.class_params() {
                if req_grad.contains(&p) && visited.insert(cid) {
                    req_grad.insert(cid);
                    topo.push(cid);
                    break;
                }
            }
        }
        topo.reverse();
        topo
    }
}

fn accum_grad(rt: &mut Runtime, graph_id: GraphId, grads: &mut Map<OpId, OpId>, nid: OpId, grad: OpId) {
    match grads.entry(nid) {
        std::collections::hash_map::Entry::Vacant(e) => {
            e.insert(grad);
        }
        std::collections::hash_map::Entry::Occupied(mut e) => {
            let eg = *e.get();
            let sum = rt.push_binary_node(graph_id, eg, grad, BOp::Add);
            e.insert(sum);
        }
    }
}

impl Runtime {
    /// Build a shape class from dim classes: rank-1 uses the dim directly,
    /// higher ranks get a `Stack`, rank-0 is `NULL`.
    pub(crate) fn shape_class(&mut self, graph_id: GraphId, dims: Vec<OpId>) -> OpId {
        match dims.len() {
            0 => OpId::NULL,
            1 => dims[0],
            _ => self.push_node(graph_id, Node::Stack { ops: dims.into_boxed_slice() }).1,
        }
    }

    /// Numeric value of a dim class, only if it is a constant. Symbolic dims
    /// return `None` — callers must not guess.
    fn graph_const_dim(&self, graph_id: GraphId, dim: OpId) -> Option<Dim> {
        self.graphs[graph_id].resolve_const(dim).and_then(|c| c.as_dim())
    }
}