The readable reverse¶

Tlaloc's reverse-mode code is readable as Kotlin source, and compiled. Every block of gradient source below was generated by the plugin or the capture API, not written by hand, and each of the three claims is pinned by a named test:

Readable. The gradient prints as val-per-op Kotlin over the :core host twins, typed with the same ranked phantom types the user writes (DTensor<Rank2<Sym, Sym>, F32>), each line annotated with the SSA node it renders.
Compiled. The printed source is plain user Kotlin over :core — it compiles standalone, with no plugin on the classpath (PrintedGradientGoldenTest, DumpGradSourceTest both compile it with an in-process K2JVMCompiler and run it).
Runnable — and bit-identical. The compiled printed source produces gradients RAW-BIT-IDENTICAL (Float.toRawBits, no epsilon) to what the engine itself computes, because the printed program IS the engine's program: the renderer prints the exact DxirFunction the interpreter evaluates and the plugin synthesises, and the :core host functions are pinned raw-bit-equal to the interpreter.

Behind all of it is one engine (AD_SINGLE_ENGINE_AUDIT.md): there is no runtime value tape. Every gradient in Tlaloc (the grad {} compiler intrinsics, the Tracer-capture API, the :nn training step) comes from DxirReverseTransform over DXIR. Tangent generated readable derivatives for Python and DiffKT ran on a runtime tape; Tlaloc generates the derivative in the Kotlin compiler and can print it.

Demo 1 — the compile-time dump (`grad {}`)¶

The user writes a grad {} lambda, exactly as always:

import io.tlaloc.autograd.grad

fun main() {
    val g = grad { x: Float -> x * x * x + 2.5f * x }
    println(g(1.75f))
    println(g(-0.5f))
}

Then asks the compiler to show the gradient it derived. The first option prints a message; the second also writes location-named .kt files:

tlaloc {
    dumpGradSource.set(true)
    dumpGradSourceDir.set(layout.buildDirectory.dir("gradients"))
}

The Gradle plugin writes each compilation's files into a subdirectory named after its source set, here build/gradients/main/, and declares it an output of the compile task, so the build cache restores it. Without the Gradle plugin the same options are -P plugin:io.tlaloc.plugin:dumpGradSource=true and -P plugin:io.tlaloc.plugin:dumpGradSourceDir=<dir>, which writes into <dir> itself.

The compiler answers with an INFO message headed by the lambda's source location (Main.kt:3:13), and — dir form — this file, verbatim as dumped by the plugin on this repository (Main_kt_3_13_grad.kt):

// Generated by DxirFunction.toKotlinSource() from function 'grad_body_grad'.
// Each `val` is one DXIR op, bound to its certified `:core` host-twin call;
// the trailing comment names the SSA node it renders.
import io.tlaloc.core.*
import io.tlaloc.core.ops.*

fun grad_body_grad(x: DTensor<ScalarShape, F32>): DTensor<ScalarShape, F32> {
    val v1: DTensor<ScalarShape, F32> = (x * x) // %1 = MUL(%0, %0)
    val v8: DTensor<ScalarShape, F32> = Tensors.f32Scalar(2.5f) // %8 = const : f32
    val v11: DTensor<ScalarShape, F32> = (v8 + v1) // %11 = ADD(%8, %1)
    val v12: DTensor<ScalarShape, F32> = (x * x) // %12 = MUL(%0, %0)
    val v14: DTensor<ScalarShape, F32> = (v11 + v12) // %14 = ADD(%11, %12)
    val v15: DTensor<ScalarShape, F32> = (v14 + v12) // %15 = ADD(%14, %12)
    return v15
}

Read it: x² + 2.5 + x² + x² = 3x² + 2.5 — the derivative of x³ + 2.5x, as the product rule actually derived it, term by term. The dump happens at the dxir level BEFORE synthesis, so this is the IDENTICAL function the plugin compiles into your bytecode — not a description of it. DumpGradSourceTest pins the whole story: the dump appears at the right location, the dumped file compiles standalone without the plugin, runs, and matches the plugin-compiled gradient raw-bit-for-bit.

A captured runtime value is a PARAMETER of the printed function¶

A grad {} body may read a runtime value declared outside it — a local val, a parameter of the enclosing function — and the lowering carries it as a TRAILING parameter of the gradient function, which the IR phase binds at the call site. The printer renders that parameter, because that is what it is:

val scale = runtimeScale(2.0f)              // not a compile-time constant
val g = grad { x: Float -> x * x * scale }

fun grad_body_grad(x: DTensor<ScalarShape, F32>, scale: DTensor<ScalarShape, F32>)
        : DTensor<ScalarShape, F32> {
    ...
}

Two properties hold, and both are pinned by CapturedRuntimeValueGradientTest.the printed gradient carries the capture as a parameter and runs standalone:

it still compiles and runs standalone, called with the captured value in that slot, raw-bit-identical to the bytecode the plugin compiled — the same certified property as Demo 1;
the return is still the user's gradient alone. The captured parameter is an input, never a differentiation target, so a one-param grad {} with three captures prints a function of four parameters returning one value — not a Pair, not a Quadruple. If the printed return type ever grew a component per capture, the returned function's own arity would have changed with it, which is the failure this pin exists to catch.

Demo 2 — an MLP loss (`:nn` capture)¶

A real two-layer MLP with a mean-squared-error loss, on the F1 training-step surface:

val model = Sequential(
    Dense(w1 /* [2,3] */, b1 /* [3] */, Activation.Relu),
    Dense(w2 /* [3,1] */, b2 /* [1] */),
)
val step = capture(model, listOf(x /* [2,2] */)) { y ->
    ((y - 0.5f) * (y - 0.5f)).mean()
}
println(step.gradSource())

step.gradSource() prints the captured gradient — the EXACT program step.run(...) executes — as this function (verbatim output on this repository; p0 is the input batch, p1..p4 are 0.w, 0.b, 1.w, 1.b in parameters order):

// Generated by DxirFunction.toKotlinSource() from function 'model_grad'.
// Each `val` is one DXIR op, bound to its certified `:core` host-twin call;
// the trailing comment names the SSA node it renders.
import io.tlaloc.core.*
import io.tlaloc.core.ops.*

fun model_grad(p0: DTensor<Rank2<Sym, Sym>, F32>, p1: DTensor<Rank2<Sym, Sym>, F32>, p2: DTensor<Rank1<Sym>, F32>, p3: DTensor<Rank2<Sym, Sym>, F32>, p4: DTensor<Rank1<Sym>, F32>): List<DTensor<*, *>> {
    val v5: DTensor<Rank2<Sym, Sym>, F32> = (p0 matmul p1) // %5 = MATMUL(%0, %1)
    val v6: DTensor<Rank2<Sym, Sym>, F32> = stretchToRank2(p2.reshape(1, 3), 2, 3) // %6 = BROADCAST(%2)
    val v7: DTensor<Rank2<Sym, Sym>, F32> = (v5 + v6) // %7 = ADD(%5, %6)
    val v8: DTensor<Rank2<Sym, Sym>, F32> = v7.relu() // %8 = RELU(%7)
    val v9: DTensor<Rank2<Sym, Sym>, F32> = (v8 matmul p3) // %9 = MATMUL(%8, %3)
    val v10: DTensor<Rank2<Sym, Sym>, F32> = stretchToRank2(p4.reshape(1, 1), 2, 1) // %10 = BROADCAST(%4)
    val v11: DTensor<Rank2<Sym, Sym>, F32> = (v9 + v10) // %11 = ADD(%9, %10)
    val v12: DTensor<Rank2<Sym, Sym>, F32> = Tensors.f32Matrix(2, 1, floatArrayOf(0.5f, 0.5f)) // %12 = const : f32[2,1]
    val v13: DTensor<Rank2<Sym, Sym>, F32> = (v11 - v12) // %13 = SUB(%11, %12)
    val v14: DTensor<Rank2<Sym, Sym>, F32> = Tensors.f32Matrix(2, 1, floatArrayOf(0.5f, 0.5f)) // %14 = const : f32[2,1]
    val v15: DTensor<Rank2<Sym, Sym>, F32> = (v11 - v14) // %15 = SUB(%11, %14)
    val v16: DTensor<Rank2<Sym, Sym>, F32> = (v13 * v15) // %16 = MUL(%13, %15)
    val v17: DTensor<ScalarShape, F32> = v16.mean() // %17 = MEAN(%16)
    val v20: DTensor<ScalarShape, F32> = Tensors.f32Scalar(0.5f) // %20 = const : f32
    val v21: DTensor<Rank2<Sym, Sym>, F32> = stretchToRank2(v20.reshape(1, 1), 2, 1) // %21 = BROADCAST(%20)
    val v22: DTensor<Rank2<Sym, Sym>, F32> = (v21 * v15) // %22 = MUL(%21, %15)
    val v23: DTensor<Rank2<Sym, Sym>, F32> = (v21 * v13) // %23 = MUL(%21, %13)
    val v26: DTensor<Rank2<Sym, Sym>, F32> = (v23 + v22) // %26 = ADD(%23, %22)
    val v27: DTensor<Rank1<Sym>, F32> = reshapeToRank1(v26.sum(0), 1) // %27 = SUM(%26)
    val v28: DTensor<Rank2<Sym, Sym>, F32> = transposePerm2(v8, 1, 0) // %28 = TRANSPOSE(%8)
    val v29: DTensor<Rank2<Sym, Sym>, F32> = transposePerm2(p3, 1, 0) // %29 = TRANSPOSE(%3)
    val v30: DTensor<Rank2<Sym, Sym>, F32> = (v26 matmul v29) // %30 = MATMUL(%26, %29)
    val v31: DTensor<Rank2<Sym, Sym>, F32> = (v28 matmul v26) // %31 = MATMUL(%28, %26)
    val v32: DTensor<Rank2<Sym, Sym>, F32> = v7.step() // %32 = STEP(%7)
    val v33: DTensor<Rank2<Sym, Sym>, F32> = (v30 * v32) // %33 = MUL(%30, %32)
    val v34: DTensor<Rank1<Sym>, F32> = reshapeToRank1(v33.sum(0), 3) // %34 = SUM(%33)
    val v35: DTensor<Rank2<Sym, Sym>, F32> = transposePerm2(p0, 1, 0) // %35 = TRANSPOSE(%0)
    val v36: DTensor<Rank2<Sym, Sym>, F32> = transposePerm2(p1, 1, 0) // %36 = TRANSPOSE(%1)
    val v37: DTensor<Rank2<Sym, Sym>, F32> = (v33 matmul v36) // %37 = MATMUL(%33, %36)
    val v38: DTensor<Rank2<Sym, Sym>, F32> = (v35 matmul v33) // %38 = MATMUL(%35, %33)
    return listOf(v17, v37, v38, v34, v31, v27)
}

It reads as the backprop a person would write on paper. Lines v5–v17 are the forward (includeForward = true — one program computes loss AND gradients). v20–v26 are the mean/square adjoint (0.5 = 1/N broadcast, product rule on the two SUB factors). v27 is d(1.b) — the bias adjoint is a column sum. v30/v31 are the textbook matmul adjoints δ·Wᵀ and hᵀ·δ; v32/v33 gate the upstream through the relu mask (step(v7)); v35–v38 repeat the pair for the first layer. The return packs (loss, d_input, d_0.w, d_0.b, d_1.w, d_1.b). PrintedGradientGoldenTest compiles this class of printed gradient standalone with K2JVMCompiler and pins it raw-bit-identical to DxirInterpreter on the same inputs; GradSourceTest pins the :nn surface itself. Running this capture (step.run) prints loss=0.5615869 with all four parameter gradients — including the exactly-zero columns the relu mask kills, visible above as v32.

The one engine¶

There is no runtime value-tape walker, so there is no second AD engine to disagree with the first:

The grad {} / grad2 {} / grad3 {} compiler intrinsics synthesise DXIR, reverse-transform it, and compile the result.
The Tracer-convenience API (grad(f), valueAndGrad, gradWithScalars) captures a tape, converts it with Tape.toDxirFunction, and hands it to the SAME DxirReverseTransform + DxirInterpreter.
The :nn training step (capture / CapturedStep) is the same route, with parameters as first-class DXIR params.

OneEngineParityTest pins the two user surfaces raw-bit-identical through the real K2 plugin; SingleEngineDeletionTest pins the tape walker's absence by Class.forName. Gradient MATH lives in exactly one place — VjpRegistry's rules, consumed by the transform — and forward math in one certified pair (interpreter arms ≡ :core host functions, pinned raw-bit-equal).

Limitations¶

Tensor grad {} compile-time dumps refuse loudly. A tensor lambda's dxir carries -1 SENTINEL dims; the renderer's ranked-literal renderings may not bake sentinel-derived values, so the dump says SKIPPED, names the sentinel reason, writes no file, and does not affect compilation. Tensor gradients still compile; only the dump refuses. The capture route (CapturedStep.gradSource()) prints tensor gradients.
Every op kind either renders or refuses by name (KotlinRenderRefusal); the renderer never prints plausible wrong source. Not rendered: ABS/RSQRT/GELU/SILU (no :core tensor function yet), rank≥3 MATMUL, GATHER/SCATTER, control flow (IF/WHILE/COARSENED), rank>4. F64 renders (Tensors.f64* constants, every digit kept).
The jvp/vjp/jacobian/hessian intrinsics do not print source. dumpGradSource covers the reverse-mode grad/valueAndGrad family.