Transform the derivative graph with shapes and saved values intact

Consider Y=XW plus a row-broadcast bias, where X has shape [B,K], W [K,M], and b [M]. An incoming cotangent G has shape [B,M], matching Y. The matrix-product rule yields an X cotangent of shape [B,K] and a W cotangent of shape [K,M]. The bias contribution must sum G over the broadcast batch dimension to return shape [M]. Simply passing the entire [B,M] G as the bias gradient violates the input's shape contract.

If an input feeds multiple downstream uses, its cotangent contributions must accumulate rather than overwrite each other. For a scalar loss x²+x, both the square path and the direct path contribute to the same input. In an SSA reverse program, each accumulation can have its own result value, while buffer lowering may implement it with an update only when alias and scheduling analysis permit that. The mathematical derivative rule and its storage implementation remain separate correctness obligations.

The W cotangent needs X, and the X cotangent needs W. If a forward operation later mutates storage that held one of these values, the reverse transform cannot silently use the changed contents in place of the required primal. It must preserve the relevant old value, choose a valid recomputation, or reject an unsupported mutation pattern. Shape metadata alone will not expose this error: the stale and correct residuals can have identical shapes while producing different derivatives.

Recomputation also interacts with effects. Repeating a pure deterministic subexpression can reconstruct a residual under its declared numerical contract; repeating a random draw without the corresponding state or key discipline may produce a different value. Repeating I/O can duplicate effects. JAX's documented JVP and VJP transforms make the linear-map direction explicit, while supported custom rules add their own contracts. A custom backward rule that returns the right shape but the wrong linear map is not repaired by successful compilation.

Check derivative shapes first, then use small exact examples, adjoint checks, or suitable finite-difference comparisons within a declared smooth domain. Test nontrivial incoming cotangents instead of always ones, since some incorrect transpose or reduction rules can accidentally agree under uniform seeds. Separately test supported aliasing, branches, and state. An IR verifier establishes structural properties; a numerical comparison gives evidence for tested cases. Neither automatically proves every custom rule, shape family, or backend lowering correct.

The network lab supplies one sigmoid primal expression for a hand-derived reverse rule. Reset, set both x=1,w=2, predict logit 1.5 and output about 0.8176, then Play or Step. Derive the sigmoid factor y(1−y), then multiply by w or x for the respective cotangent with seed one. Reset and set both x=1,w=1 for comparison. The instrument displays the primal network, not an AD tape, saved-tensor analysis, or a compiler-generated backward graph.