34. Differentiating and Training Through the Dynamic Compiler

A forward tensor compiler is not yet a training compiler. Backward needs graph accumulation, shape-aware derivative rules, and a runtime that observes updated parameters without baking their values or batch sizes into the executable. This chapter adds those pieces to tensorjit and runs an actual generated forward/backward program.

.venv/bin/python -m examples.train_tensorjit
.venv/bin/python -m pytest tests/test_tensorjit.py -q

The model is deliberately small enough to validate against a known answer. Data are generated by target = X*true_W + true_bias. The compiler differentiates squared error with respect to W and bias, then returns loss and both gradients. NumPy performs the SGD update outside the compiled graph. This is generated-kernel training with a host optimizer, not a claim that the entire training loop has been captured or that the model exercises large-scale language-model behavior.

34.1 Differentiate before scheduling

Reverse mode operates on the pure tensor graph, where matrix multiplication, reduction, and broadcasting still carry their mathematical meaning. It does not differentiate generated stores, allocation choices, or C loop syntax. Once backward is a graph, the same materialization, fusion, shape-binding, and runtime machinery compiles it.

Seed a scalar loss with gradient 1 and visit graph nodes in reverse topological order. Each local rule contributes gradients to its sources. If a source receives contributions from several consumers, add them. For loss=sum(x*x), the same x appears twice as a multiplication source; both contributions are required. Replacing the earlier contribution instead of accumulating it halves the derivative.

A topological traversal follows node identity and visits each shared node once. The adjoint graph can still refer to the original forward values. The combined output graph of loss and requested gradients lets the scheduler identify common forward/backward producers and materialize shared work. This is graph sharing, not checkpointing: recomputation policies are a later extension.

34.2 Reverse broadcasting with reductions

For Z=X+b with X:(B,N) and b:(N), an upstream dZ:(B,N) contributes directly to X. The gradient for b must sum over B. Broadcasting repeated b's value in each row, so reverse mode accumulates the contributions of all those logical uses.

General unbroadcasting first removes extra leading axes by summing axis zero. Then, for every original singleton dimension, reduce the corresponding gradient axis and reinsert the singleton. The final gradient shape must equal the source shape exactly. Do not merely rely on a later broadcast to hide a wrong derivative shape.

For X:(M,1) broadcast to (M,N), dX has shape (M,1), not (M,). Removing the reduced axis without reinserting it would silently reinterpret which axis repeats in subsequent arithmetic. The test with an explicit (2,1) broadcast to (2,4) checks both values and shape.

34.3 Reduction and movement derivatives

A sum removes one axis. Backward inserts that axis as a singleton, then broadcasts to the source shape. A transpose reverses its matrix axes, and its derivative applies the same transpose to the upstream gradient. Inserting a singleton reverses to a reduction over that axis; explicit broadcast reverses through unbroadcasting.

These movement nodes need no standalone buffers when they have one consumer and can be expressed as index maps. Their semantics remain visible during differentiation, then their index expressions are inlined during code generation. This is the same separation of mathematical and storage concerns used in the original compiler, now exercised with runtime dimension arguments.

34.4 Matrix derivatives preserve symbolic dimensions

For Y=XW, with X:(B,K) and W:(K,N), the upstream gradient has shape (B,N):

dX = dY @ transpose(W)       # (B,N) @ (N,K) -> (B,K)
dW = transpose(X) @ dY       # (K,B) @ (B,N) -> (K,N)

The reduction extent B in dW is the original batch symbol, not the batch size of an example call. Compile once and bind B from each invocation's inputs. If differentiation accidentally reads a concrete batch length while constructing the graph, changing batch length can produce stale loops or a new graph per shape. The training test alternates several lengths and checks one library build.

ReLU contributes dY * (X>0), with zero derivative chosen at zero. The implementation uses a float32 mask node; it is not a differentiable general comparison primitive. Higher derivatives through that mask are explicitly unsupported. Numerical derivative tests use finite values away from ReLU's nondifferentiable points.

"""Reverse-mode differentiation of the pure dynamic tensor graph.

Differentiation precedes fusion and code generation. The shape expressions in
backward retain the original symbols; no concrete runtime extent is captured.
"""
from tensorjit.graph import Expr, constant, toposort


def unbroadcast(gradient, shape):
  while len(gradient.shape) > len(shape): gradient = gradient.sum(0)
  for axis, extent in enumerate(shape):
    if extent == 1 and gradient.shape[axis] != 1:
      gradient = gradient.sum(axis).expand_dims(axis)
  if gradient.shape != shape: raise ValueError("gradient shape does not match source")
  return gradient


def gradients(loss, wrt):
  if not isinstance(loss, Expr) or loss.shape != (): raise ValueError("loss must be a scalar expression")
  wrt = tuple(wrt)
  if any(not isinstance(node, Expr) or node.op != "input" for node in wrt): raise ValueError("differentiate with respect to input nodes")
  order = toposort((loss,))
  accumulated = {loss: constant(1)}
  for node in reversed(order):
    if node not in accumulated: continue
    gradient, sources = accumulated[node], node.sources
    if node.op in ("input", "constant"): continue
    if node.op == "add": contributions = (unbroadcast(gradient, sources[0].shape), unbroadcast(gradient, sources[1].shape))
    elif node.op == "mul":
      contributions = (unbroadcast(gradient*sources[1], sources[0].shape), unbroadcast(gradient*sources[0], sources[1].shape))
    elif node.op == "relu": contributions = (gradient*Expr("positive", sources[0].shape, sources),)
    elif node.op == "transpose": contributions = (gradient.transpose(),)
    elif node.op == "sum": contributions = (gradient.expand_dims(node.arg).broadcast_to(sources[0].shape),)
    elif node.op == "expand_dims": contributions = (gradient.sum(node.arg),)
    elif node.op == "broadcast": contributions = (unbroadcast(gradient, sources[0].shape),)
    elif node.op == "matmul": contributions = (gradient @ sources[1].transpose(), sources[0].transpose() @ gradient)
    elif node.op == "positive":
      raise ValueError("higher derivatives through the ReLU mask are unsupported")
    else: raise ValueError(f"missing derivative rule: {node.op}")
    for source, contribution in zip(sources, contributions):
      if contribution.shape != source.shape: raise ValueError("derivative rule produced wrong shape")
      accumulated[source] = accumulated[source]+contribution if source in accumulated else contribution
  return tuple(accumulated[node] if node in accumulated else constant(0).broadcast_to(node.shape) for node in wrt)

34.5 Shape information can outlive data dependencies

The derivative of sum(X) is all ones. Its generated graph may no longer read X at all, but the shape of the ones still contains X's symbolic dimensions. Removing a data dependency cannot justify inventing a dimension binding.

Ordinary calls infer symbols from actual input array metadata. For shape-only graphs, use compiled.run_with_bindings({"N": length}, **inputs). Explicit bindings are checked for known names, integer type, and declared bounds. If an input also establishes N, the explicit binding must match it. This supports constant gradients without forcing a dummy data read or bypassing the shape contract.

An absent differentiated input gets a zero expression broadcast to its declared shape. If that shape contains an otherwise unbound symbol, supply the same checked metadata. No unrequested buffer value is read just to obtain an arbitrary default dimension.

Experiment: differentiate a sum of 2*(X+row) with respect to X, row, and a disconnected input. The gradients are all twos, repeated-row counts times two, and zeros, respectively. Compile the gradients alone and vary N through explicit bindings. Then provide a conflicting input N and verify that dispatch fails before native work.

34.6 A reusable forward/backward executable

The training loss is:

E = XW + bias - target
L = sum(E*E) * scale

Scale is a scalar float32 input equal to 1/(batch_size*output_width), keeping the mean-loss convention consistent across batch sizes. Encoding that changing normalization as data avoids rebuilding the graph. Alternatively, a richer symbolic compiler could expose dimension values as scalar graph expressions; this implementation keeps the boundary explicit.

Initialize W and bias to zero, generate deterministic batches with sizes 8, 16, and 32 in rotation, execute (loss,dW,dbias), and update the actual input arrays with learning rate 0.15. Since the runtime reads parameters on every invocation, the next call observes the updated values. It does not cache buffer contents as constants.

"""Train a linear model through generated dynamic forward/backward kernels.

Run: python -m examples.train_tensorjit. Requires gcc; no PyTorch dependency.
Changing minibatch length reuses one compiled library. NumPy performs SGD updates.
"""
import numpy as np
from tensorjit import CompiledGraph, Dim, gradients, input


def train(steps=100, seed=7):
  if steps < 1: raise ValueError("steps must be positive")
  batch = Dim("B", 1, 64)
  x, y = input("x", (batch, 3)), input("y", (batch, 2))
  w, bias, scale = input("w", (3, 2)), input("bias", (2,)), input("scale", ())
  error = x @ w + bias - y
  loss = (error*error).sum(0).sum(0)*scale
  dw, db = gradients(loss, (w, bias))
  compiled = CompiledGraph((loss, dw, db))
  rng = np.random.default_rng(seed)
  true_w, true_bias = rng.normal(size=(3, 2)).astype(np.float32), rng.normal(size=2).astype(np.float32)
  weights, biases = np.zeros((3, 2), dtype=np.float32), np.zeros(2, dtype=np.float32)
  history = []
  for step in range(steps):
    length = (8, 16, 32)[step % 3]
    data = rng.normal(size=(length, 3)).astype(np.float32)
    target = data @ true_w + true_bias
    actual_loss, grad_w, grad_b = compiled(x=data, y=target, w=weights, bias=biases, scale=np.array(1/(length*2), dtype=np.float32))
    history.append(float(actual_loss))
    weights -= np.float32(0.15)*grad_w
    biases -= np.float32(0.15)*grad_b
  return {"initial_loss": history[0], "final_loss": history[-1], "parameter_max_error": float(np.abs(weights-true_w).max()),
          "bias_max_error": float(np.abs(biases-true_bias).max()), "compiled_variants": compiled.compilations,
          "calls": compiled.calls, "kernels": len(compiled.plan.kernels)}


def main():
  result = train()
  for key, value in result.items(): print(f"{key}: {value}")
  assert result["parameter_max_error"] < 1e-3 and result["bias_max_error"] < 1e-3


if __name__ == "__main__": main()

A recorded 100-step CPU run with seed 7 produced:

These are different minibatches, so the loss sequence need not be monotonic. The stronger result is recovery of the known generating parameters, not simply a favorable first/last loss comparison. The dataset is noiseless and the model is linear; this is an integration test, not a generalization claim or evidence of neural-network training quality.

34.7 Validate backward independently of its own forward

Agreement between generated forward and generated backward can conceal a shared mistake. The tests compute the loss with an independent double-precision NumPy expression, perturb each input along a random direction, and compare its central finite difference with the inner product of the compiled gradient and that direction. Perturbations are evaluated in the double-precision oracle rather than being rounded into the compiled float32 inputs.

Separate tests cover repeated uses, leading-axis broadcasting, singleton-axis broadcasting, transpose, reduction, disconnected inputs, dynamic batch sizes, and explicit shape-only metadata. The training test checks parameter recovery and executable reuse after 100 parameter updates. A gradient test should fail on either wrong values or wrong shapes; an accidentally broadcast-compatible answer is not accepted.

Next implementation: add stable softmax and normalization with independently derived gradients; checkpoint a pure nonlinear block; implement Adam with explicit state; and lower this dynamic forward/backward path to a GPU backend. Keep numerical policy, alias ownership, and state updates specified before optimizing. This chapter establishes the CPU training path that those extensions must preserve.