Module vjp 

Reverse-mode differentiation of a tensor program, as a program.

backward reads a forward program’s tensor statements (the typed specs the front end lifts, crate::spatial::passes::recognize_stages) and writes the statements that compute the gradient of a scalar loss with respect to any tensor the forward reads. The backward pass is ORDINARY LOGOS — tensor sugar where the language has the form (matmul, transpose, plus, hadamard, scaled by), row-major loops over declared tensors where it does not — so the same front end, fusion and every backend take it, and the VM runs it in f64 as the oracle a gradient is checked against.

The loss is Σ out ⊙ cotangent for a declared cotangent tensor of the output’s shape: the seed of reverse mode is that tensor, and every gradient follows from it. A tensor read by several stages sums its gradients.

Every op kind has a rule or a refusal that names it (vjp_stage, an exhaustive match over TensorOpSpec). A refusal is either VjpRefusal::NoGradient — the op is piecewise constant in its inputs (argmax) — or VjpRefusal::NotDerived, an op whose rule is not written yet: a training program that reaches it is refused by name, never differentiated as if the op were absent.

Structs

Backward

A derived backward pass.

Enums

VjpRefusal

Why a stage has no backward rule.

Functions

backward

The backward pass of forward (a program whose ## Main holds tensor statements) for the loss Σ output ⊙ cotangent, giving the gradient of every tensor in wrt. cotangent must be a declared tensor of output’s shape. A tensor in wrt the loss does not depend on gets a zero gradient.

backward_of

backward over a program the front end has already lifted.