readable-gradients — read the derivative the compiler wrote¶

What it shows. Every autodiff framework hands you a number. Tlaloc hands you the program: the Kotlin source of the gradient it derived, which you can read, review in a pull request, compile on its own, and run.

This project does exactly that, three times over, and checks that all three answers are the same one:

	where it comes from	what it proves
1. compiled `grad { }`	the K2 plugin rewrote the lambda into synthesized bytecode at compile time	the fast path — no tape, no graph object, no framework at runtime
2. the printed source	the same derivation, dumped as Kotlin by the `dumpGradSourceDir` plugin flag, compiled by a source set with no Tlaloc plugin on its classpath, called here by reflection	what you read is what ran, and it is ordinary Kotlin
3. a central difference	a hand-written `kotlin.math` transcription of the same formula that never touches Tlaloc	the derivation is correct, checked by something that shares no code with it

1 and 2 are the same program, so they agree raw-bit-identically — the example asserts Float.toRawBits() equality, with no epsilon. 3 is numerical, so it agrees to finite-difference tolerance (observed: 5.5e-08).

The function¶

Deliberately not a parabola. A "damped swish":

swish(x)    = x · σ(x)              a smooth, non-monotonic gate
softplus(x) = log(1 + eˣ)           a smooth, positive magnitude
f(x)        = swish(x) / √(1 + softplus(x))

In Kotlin, inside the block the compiler differentiates:

val g = grad { x: Float ->
    (x * x.sigmoid()) / (1.0f + (1.0f + x.exp()).log()).sqrt()
}

By hand, f'(x) is a quotient rule whose numerator needs the product rule through a sigmoid, over a denominator that chains sqrt ∘ log ∘ exp:

f'(x) = [σ(x) + x·σ(x)(1 − σ(x))] / √(1+s)  −  x·σ(x)² / (2·(1+s)^1.5),   s = softplus(x)

Four rules, two shared subexpressions, and an easy sign to get wrong. Nobody writes that by hand. Below is the machine's version of it — and unlike a number, you can check it.

Running it¶

A standalone Gradle project that resolves Tlaloc from mavenLocal, exactly as your own project would, so publish first:

# from the Tlaloc repo root
./gradlew publishToMavenLocal
./gradlew -p examples/readable-gradients run

No GPU, no native library, no Python — this one runs anywhere a JVM does.

Expected output¶

Verbatim, from this machine (GB10 / aarch64, JDK 25, Kotlin 2.4.20):

Tlaloc readable gradients
  f(x) = x*sigmoid(x) / sqrt(1 + log(1 + exp(x)))

--- the gradient, as the compiler printed it during the build ---
--- Main_kt_72_13_grad.kt (generated; not in src/) ---
// 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.sigmoid() // %1 = SIGMOID(%0)
    val v2: DTensor<ScalarShape, F32> = (x * v1) // %2 = MUL(%0, %1)
    val v3: DTensor<ScalarShape, F32> = Tensors.f32Scalar(1.0f) // %3 = const : f32
    val v5: DTensor<ScalarShape, F32> = x.exp() // %5 = EXP(%0)
    val v6: DTensor<ScalarShape, F32> = (v3 + v5) // %6 = ADD(%3, %5)
    val v7: DTensor<ScalarShape, F32> = v6.log() // %7 = LOG(%6)
    val v8: DTensor<ScalarShape, F32> = (v3 + v7) // %8 = ADD(%3, %7)
    val v9: DTensor<ScalarShape, F32> = v8.sqrt() // %9 = SQRT(%8)
    val v12: DTensor<ScalarShape, F32> = (v3 / v9) // %12 = DIV(%3, %9)
    val v13: DTensor<ScalarShape, F32> = (v9 * v9) // %13 = MUL(%9, %9)
    val v14: DTensor<ScalarShape, F32> = (v2 / v13) // %14 = DIV(%2, %13)
    val v16: DTensor<ScalarShape, F32> = v14.neg() // %16 = NEG(%14)
    val v17: DTensor<ScalarShape, F32> = Tensors.f32Scalar(2.0f) // %17 = const : f32
    val v19: DTensor<ScalarShape, F32> = (v17 * v9) // %19 = MUL(%17, %9)
    val v20: DTensor<ScalarShape, F32> = (v16 / v19) // %20 = DIV(%16, %19)
    val v21: DTensor<ScalarShape, F32> = (v20 / v6) // %21 = DIV(%20, %6)
    val v23: DTensor<ScalarShape, F32> = (v21 * v5) // %23 = MUL(%21, %5)
    val v24: DTensor<ScalarShape, F32> = (v12 * v1) // %24 = MUL(%12, %1)
    val v25: DTensor<ScalarShape, F32> = (v12 * x) // %25 = MUL(%12, %0)
    val v26: DTensor<ScalarShape, F32> = (v23 + v24) // %26 = ADD(%23, %24)
    val v29: DTensor<ScalarShape, F32> = (v3 - v1) // %29 = SUB(%3, %1)
    val v30: DTensor<ScalarShape, F32> = (v1 * v29) // %30 = MUL(%1, %29)
    val v31: DTensor<ScalarShape, F32> = (v25 * v30) // %31 = MUL(%25, %30)
    val v32: DTensor<ScalarShape, F32> = (v26 + v31) // %32 = ADD(%26, %31)
    return v32
}

--- three answers that must agree ---
       x     compiled grad{}      printed source  central difference
   -3.00         -0.08289663         -0.08289663         -0.08289664
   -1.00          0.08714630          0.08714630          0.08714630
   -0.25          0.31185716          0.31185716          0.31185718
    0.00          0.38425776          0.38425776          0.38425776
    0.50          0.49173251          0.49173251          0.49173249
    1.75          0.50086808          0.50086808          0.50086807
    4.00          0.29834056          0.29834056          0.29834050

compiled == printed: raw-bit identical at every point
compiled vs finite difference: max |gap| = 5.525e-08
readable-gradients OK

That block between the dashes is not a log line the example formatted. It is a file — build/gradients/Main_kt_72_13_grad.kt, named after the source location of the grad { } call it came from — that the Kotlin compiler wrote while compiling Main.kt, and that a second compilation turned into the .class file the third column was computed by.

Reading the gradient¶

It is worth spending sixty seconds on it, because this is the part no other framework lets you do.

The forward, v1–v9 — only the pieces the adjoints actually need:


`v1 = σ(x)`, `v2 = x·σ(x)`	the numerator, `swish`
`v5 = eˣ`, `v6 = 1+eˣ`, `v7 = log(1+eˣ)`	`softplus`, built up
`v8 = 1+s`, `v9 = √(1+s)`	the denominator, call it `d`

Note what is missing: v2 / v9, the value of f itself, is never formed. grad { } asks for the derivative only, so the final division is dead code and the transform drops it. (valueAndGrad { } keeps it.)

The reverse, v12–v32, seeded with an upstream of 1:


`v12 = 1/d`	quotient rule, ∂/∂numerator
`v13 = d²`, `v14 = num/d²`, `v16 = −num/d²`	quotient rule, ∂/∂denominator
`v17 = 2`, `v19 = 2d`, `v20 = v16/(2d)`	the `sqrt` rule, `d√u = 1/(2√u)`
`v21 = v20/(1+eˣ)`	the `log` rule, `d log u = 1/u`
`v23 = v21·eˣ`	the `exp` rule — the denominator's contribution to `dx`
`v24 = (1/d)·σ(x)`	product rule, first term
`v25 = (1/d)·x`, `v29 = 1−σ`, `v30 = σ(1−σ)`, `v31 = v25·v30`	product rule, second term, through the sigmoid rule `σ′ = σ(1−σ)`
`v26`, `v32`	the two accumulations that add the paths back together

Two things a careful reader will notice, and both are true:

v23/v21 compute eˣ/(1+eˣ), which is σ(x) — already sitting in v1. The transform applies the log and exp rules literally rather than recognising the identity. Tlaloc does not pretend to be a computer algebra system that will out-clever you; it shows you what it did, and you can see the missed simplification for yourself. That is the trade this design makes.
v3 (the constant 1.0) is used four times. The primal lowered two separate const 1.0 nodes (1.0f + appears twice in the source); the simplify pass merged them before the dump. The gradient you read is the simplified one — the same one that became bytecode.

The gaps in the numbering are not accidental either: %4, %10, %11, … are nodes that existed and were then eliminated — merged as duplicates, or dropped as dead. The renderer prints the surviving SSA names rather than renumbering, so the file stays diffable against the IR it came from.

How the build does it¶

The whole mechanism is nine lines of build.gradle.kts, and it is worth reading because the second half is the interesting one:

// 1. Ask the plugin to dump every gradient it synthesises, as Kotlin.
tasks.named<KotlinCompile>("compileKotlin") {
    compilerOptions.freeCompilerArgs.addAll(
        "-P", "plugin:io.tlaloc.plugin:dumpGradSourceDir=${gradientSourceDir.asFile.absolutePath}",
    )
    outputs.dir(gradientSourceDir)
}

// 2. Compile the dump in a source set that does NOT get the plugin.
sourceSets { create("printed") }
kotlin.sourceSets.named("printed") { kotlin.srcDir(gradientSourceDir) }
dependencies { "printedImplementation"("io.github.pedronahum:tlaloc-core:0.1.0-alpha03") }

Step 2 is the claim under test. printedImplementation gets :core — the host tensor library — and nothing else. The compiler plugin is declared on kotlinCompilerPluginClasspathMain, which the printed compilation does not inherit, so the Tlaloc plugin is not present for that compilation. If the printed gradient compiles there, it is plain Kotlin; if it then produces the same bits, the source you were shown is the code that ran.

While you are at it, watch the build log. The plugin also reports the primal it lowered, as a compiler warning at the grad { } call site:

w: file:///…/Main.kt:72:13 Tlaloc lowered lambda to dxir:
fn grad_body(%0: f32) -> f32 {
  %1 = sigmoid(%0) : f32
  %2 = mul(%0, %1) : f32
  …

Forward IR in the log, reverse source on disk. Nothing about the derivation is hidden from you.

What the compiler refuses¶

The other half of "the compiler knows what you meant" is what it will not compile. src/main/kotlin/Refusals.kt.disabled holds two gradients that cannot exist. The file is named .disabled so Gradle never compiles it and this project stays green; rename it to see the errors:

cd examples/readable-gradients
mv src/main/kotlin/Refusals.kt.disabled src/main/kotlin/Refusals.kt
./gradlew compileKotlin        # fails
mv src/main/kotlin/Refusals.kt src/main/kotlin/Refusals.kt.disabled

(a) An axis error. A contraction between an operand indexed [Batch, SeqLen] and one indexed [Hidden, Hidden] — there is no axis to sum over. Verbatim, from the run above:

e: Refusals.kt:40:24 Tlaloc named-index mismatch: contract operands share no named axis: lhs=[Batch, SeqLen] rhs=[Hidden, Hidden]

Two independent layers catch this. The axis names are phantom type parameters, so Kotlin's own type checker rejects the call as well (it follows with Cannot infer type for type parameter 'NameN' and an argument-type mismatch naming the full Named<…> types). Tlaloc's diagnostic is the one that says which axes, by name, in one line.

(b) A differentiability error. grad { } over a body that returns a tensor. There is no single gradient of a tensor-valued function:

e: Refusals.kt:63:24 Tlaloc cannot differentiate this body: DxirReverseTransform v1 requires a scalar return (got f32[-1,-1]) — tensor-returning primals need seedAsParam (the vjp pullback form)

Kotlin types this one happily — grad(f: (A) -> R) leaves R unconstrained — so the refusal comes from Tlaloc running the real reverse transform at check time, in the compiler, on the body you wrote. The message names the transform, the offending type, and the surface you actually wanted (vjp, the seeded pullback). The -1s are Tlaloc's symbolic extents: the dims are not known until you call the function, which is precisely why the refusal is about the return kind and not about sizes.

Neither of these is a runtime exception, a NaN, or a silently wrong number an hour into a training run. Both are red squiggles in the IDE.

Limitations¶

The dump covers scalar grad { } bodies. A tensor grad { } lambda still compiles and still differentiates correctly, but its compile-time dump refuses loudly (a SKIPPED message naming the reason) rather than printing source whose shape literals would be derived from symbolic -1 extents.
Every op either renders or refuses by name, instead of emitting plausible-looking wrong source. Not rendered: ABS/RSQRT/GELU/SILU, rank ≥ 3 matmul, gather/scatter, control flow, and F64.
The jvp / vjp / jacobian / hessian intrinsics do not print source.

The full statement of what is certified, and the named tests that pin it (including a golden test that compiles printed gradients standalone and checks them raw-bit-identical against the interpreter), is in docs/READABLE_REVERSE.md. This example is that document, made into something you can run.

What to read in the source¶

Look at	For
`src/main/kotlin/Main.kt`	the `grad { }` call, the independent `reference()` used by the finite difference, and the reflection that loads the printed gradient
`src/main/kotlin/Refusals.kt.disabled`	the two gradients the compiler refuses, and why
`build.gradle.kts`	the `dumpGradSourceDir` flag, and the plugin-free `printed` source set that is the whole point
`build/gradients/` (after a build)	the generated gradient itself