Tlaloc

Compile-time automatic differentiation and GPU serving for Kotlin

Transformations | Training | Serving | Install | Examples | Changelog | Docs

Tlaloc is a Kotlin library for differentiable array programs. A K2 compiler plugin turns grad { } into gradient code at compile time. Tensors carry their rank, dtype and axis names in the Kotlin type. Programs lower to StableHLO and run on NVIDIA GPUs through PJRT or IREE, on Apple-silicon GPUs through MetalHLO (below), on the Cerebras Wafer-Scale Engine simulator (work in progress, below), or on the CPU through a reference interpreter.

import io.tlaloc.autograd.grad
import io.tlaloc.core.*
import io.tlaloc.core.ops.*

fun main() {
    // Rewritten at compile time: d/dA sum(A·A) = 1·Aᵀ + Aᵀ·1
    val g = grad { a: DTensor<Rank2<Sym, Sym>, F32> -> (a matmul a).sum().toFloat() }

    val a = Tensors.f32Matrix<Sym, Sym>(2, 2, floatArrayOf(1f, 2f, 3f, 4f))
    println(g(a).hostF32().toList())   // [7.0, 11.0, 9.0, 13.0]
}

grad { f } returns a plain Kotlin function. There is no tape and no runtime tracing. A model compiled for serving is a directory of StableHLO plus weights; serving it needs a PJRT plugin and a GPU driver, not a JVM.

This is an alpha, 0.1.0-alpha02. APIs can change between alphas without a deprecation cycle (COMPATIBILITY.md).

• Transformations
• Typed tensors
• Training
• Serving
• Supported platforms
• Installation
• Sharp edges
• Examples
• Documentation

Each transformation takes a lambda and returns a function. The compiler plugin builds the derivative when your code compiles.

val df = grad { x: Float -> x * x * x }
df(2f)                                    // 12.0

val dfdxy = grad2 { x: Float, y: Float -> x * y + x }

grad2 and grad3 take two and three arguments. valueAndGrad returns the value too.

val f   = jvp { x: Float -> x * x * x }
f(2f, 1f)                                 // 12.0, the directional derivative

val j = jacobian { x: DTensor<Rank1<Sym>, F32> -> x * x }                   // diag(2x)
val h = hessian  { x: DTensor<Rank1<Sym>, F32> -> (x * x).sum().toFloat() } // 2·I

vjp and jacobianReverse are also available, and the transformations nest.

// one gradient of w1 per example; w1 and w2 are shared by every example
val g1 = vmap2(batchAxis(Batch), Batched, Batched) { x: X, y: Y ->
    grad { w: W1 -> crossEntropyLoss((x matmul w).tanh() matmul w2, y).toFloat() }(w1)
}
g1(xs, ys)  // DTensor<Rank3<Named<Batch, Sym>, Named<Feature, Sym>, Named<Hidden, Sym>>, F32>

vmap writes the batched function at compile time. The batch axis is a named axis in the type, so a batch of another axis does not compile. grad { vmap { } } differentiates a batched loss; vmap2 shares or batches each of two arguments. examples/per-example-gradients computes per-example gradient norms. → design/vmap.md

The same programs run in F64, in double precision end to end: literals, captured Doubles and inputs are never rounded to F32, and F32 and F64 never mix implicitly (mixing them is a compile error at the call).

val g = grad { k: DTensor<Rank2<Sym, Sym>, F64> -> k.logDetSpd().toDouble() }   // K⁻¹
val h = hessian { k: DTensor<Rank2<Sym, Sym>, F64> -> k.logDetSpd().toDouble() } // matches JAX x64 to 1e-11

examples/gaussian-process fits a Gaussian process in F64.

if, when and for loops inside the lambda are differentiated. examples/differentiable-physics differentiates a 38-step Euler integrator. customVjp, customJvp and customVjpJvp supply your own derivative for a function.

The lambda can use values declared outside it. Compile-time constants are folded in; other values become parameters of the generated function. A lambda the plugin cannot differentiate is a compile error that names the reason.

tlaloc {
    dumpGradSource.set(true)                                   // print it
    dumpGradSourceDir.set(layout.buildDirectory.dir("grads"))  // one .kt per lambda
}

The plugin prints the derivative as Kotlin source. That source compiles without the plugin and returns the same bits as the compiled gradient. examples/readable-gradients checks this at 7 points. → READABLE_REVERSE.md

val activations: DTensor<Rank2<Named<Batch, Sym>, Named<SeqLen, Sym>>, F32> = ...
val weights:     DTensor<Rank2<Named<SeqLen, Sym>, Named<Hidden, Sym>>, F32> = ...

val hidden = activations contract weights   // compiles: both have SeqLen

With a Hidden × Hidden weight the call does not compile. Shape and differentiability errors are reported at the call's file, line and column. → examples/named-indices

An axis can carry an upper bound (experimental):

object MaxSeq : DimBound(4096)

val x = Tensors.f32Zeros<Named<SeqLen, Bounded<MaxSeq>>, Sym>(5000, 64)  // compile error: 5000 > 4096

A program over bounded axes runs at any size up to the bound, and exports as one program per bucket, so a server or a training loop compiles once per bucket instead of once per length. → design/bounded-dims.md

:nn has immutable layers and pure optimizers: a training step returns a new model and a new optimizer state.

val keys = RandomKey.fromSeed(7).split(2)
val model0 = Sequential(Dense(2, 16, keys[0]), ReluLayer, Dense(16, 1, keys[1]))

val step = capture(model0, listOf(x), name = "mlp") { prediction ->
    val residual = prediction - prediction.constant<Shape>(targets, intArrayOf(n, 1))
    (residual * residual).mean()
}

val optimizer = Adam(learningRate = 0.02f)
var model = model0
var state = optimizer.initialState()
repeat(60) {
    val out = step.run(model, listOf(x))   // loss and gradients
    val (nextModel, nextState) = optimizer.step(model, out.gradients, state)
    model = nextModel; state = nextState
}

saveCheckpoint(path, model, optimizer, state)   // one safetensors file

Kotlin reads a Hugging Face checkpoint and writes a serving artifact: StableHLO prefill and decode programs, the weights as safetensors, and a manifest. Three runtimes load it without a JVM, and one inside the JVM. All four compile it with a PJRT plugin.

Models. Each one generates the same greedy token ids as Hugging Face transformers.

Muse Glimmer runs contexts up to 32,768 tokens and several sequences at once. A 31,744-token prompt prefills in 62 s. Int8 weights keep its fixture answers and change perplexity from 5.610 to 5.602.

Serving features: paged attention; sliding-window layers whose KV cache is bounded by the window; chunked and batched prefill; KV cache updated in place; continuous batching of decode steps; CUDA shared-memory inputs and outputs; dynamic batching for stateless models.

examples/triton-llm exports Qwen3-0.6B, starts Triton and streams a chat answer in one script.

Programs with bounded axes (experimental) export as one StableHLO program per bucket and a manifest. The Python runtime and Triton pad each request to the smallest bucket that holds it: 200 requests of lengths 1..512 compile 6 programs instead of 163. → SERVING_ARCHITECTURE.md

NVIDIA results come from one machine, an NVIDIA GB10 (Blackwell, aarch64); Metal results from an Apple M-series Mac. CI builds and tests the library on the three operating systems in the table, with no GPU. → CAPABILITIES.md lists every capability and the test that certifies it.

Tlaloc runs on the GPU of an Apple-silicon Mac through MetalHLO's PJRT plugin: the same StableHLO programs, the same PjrtSession API, PjrtTarget.Metal as the target. The whole runtime-pjrt device suite — training steps, linear algebra, vmap, RNG, paged attention, MoE, FP8 KV caches, Qwen greedy decoding — passes on Metal (142 passed, 0 failed).

git clone https://github.com/pedronahum/MetalHLO.git   # next to your tlaloc checkout
cd MetalHLO && swift build -c release --product PJRTMetalHLO

PjrtSession(plugin = PjrtBinaries.requireMetalPlugin("app"), target = PjrtTarget.Metal).use { session ->
    val out = session.runOn(fn, inputs)
}

The plugin is found next to the checkout, or through METALHLO_PLUGIN_PATH (GETTING_STARTED.md). Run the device suite on Metal with TLALOC_TEST_PJRT_TARGET=metal ./gradlew :runtime-pjrt:jvmTest. Metal support is on main, not yet in 0.1.0-alpha02: build from source. Metal has no double precision, so f64 programs compute in f32.

Tlaloc also runs on the Cerebras Wafer-Scale Engine (WSE), today on the Cerebras SDK's simulator, through the PJRT plugin of CerebrasStableHLO, with PjrtTarget.Cerebras as the target. The plugin lowers each StableHLO program to CSL through CerebrasHLO, its compiler for the wafer, and compiles it with cslc.

Work in progress. CerebrasStableHLO and CerebrasHLO are not public yet; they will be made public soon. Until then the links above do not resolve.

Status of the runtime-pjrt device suite on the simulator: at the last full run, 134 tests passed, 68 were skipped and 11 failed (Metal: 147 passed, 66 skipped). Two of the five failing classes, GatedDeltaOps and BlockwisePagedAttention, have since been fixed by spreading a computation that is too big for one processing element over a group of them. The other three (MoeExperts, Qwen35MoeParity, CausalLmTraining) still fail with typed compile errors about scale; no wrong result has been measured. The simulator is slow: a test class takes minutes, and the biggest take hours.

# CerebrasStableHLO and CerebrasHLO checked out next to your tlaloc checkout,
# plus the Cerebras SDK in a Lima VM named cs_sdk
cd CerebrasStableHLO && cmake -G Ninja -B build \
  -DMLIR_DIR=/opt/homebrew/opt/llvm/lib/cmake/mlir \
  -DLLVM_DIR=/opt/homebrew/opt/llvm/lib/cmake/llvm && cmake --build build

PjrtSession(plugin = PjrtBinaries.requireCerebrasPlugin("app"), target = PjrtTarget.Cerebras).use { session ->
    val out = session.runOn(fn, inputs)
}

The plugin is found next to the checkout, or through CEREBRASSTABLEHLO_PLUGIN_PATH (GETTING_STARTED.md). Set CEREBRASSTABLEHLO_BACKEND=simulator to run on the wafer simulator; without it the plugin runs a host interpreter. Run one device test class with CEREBRASSTABLEHLO_BACKEND=simulator TLALOC_TEST_PJRT_TARGET=cerebras ./gradlew :runtime-pjrt:jvmTest --tests <class>. Cerebras support is on main, not in 0.1.0-alpha02. The WSE has no double precision, so f64 programs compute in f32.

From Maven Central. The Gradle plugin is also on Maven Central, not on the Gradle Plugin Portal, so pluginManagement needs mavenCentral().

// settings.gradle.kts
pluginManagement { repositories { mavenCentral(); gradlePluginPortal() } }
dependencyResolutionManagement { repositories { mavenCentral() } }

// build.gradle.kts
import org.jetbrains.kotlin.gradle.dsl.JvmTarget

plugins {
    kotlin("jvm") version "2.4.20"
    id("io.github.pedronahum.tlaloc") version "0.1.0-alpha02"  // makes `grad { }` compile-time
    application
}
kotlin {
    jvmToolchain(25)                                     // to BUILD
    compilerOptions { jvmTarget.set(JvmTarget.JVM_21) }  // to RUN on JDK 21
}
java {                                                   // Java tasks match Kotlin's target
    sourceCompatibility = JavaVersion.VERSION_21
    targetCompatibility = JavaVersion.VERSION_21
}

dependencies {
    implementation(platform("io.github.pedronahum:tlaloc-bom:0.1.0-alpha02"))
    implementation("io.github.pedronahum:tlaloc-autograd")
    implementation("io.github.pedronahum:tlaloc-nn")           // layers + optimizers
    implementation("io.github.pedronahum:tlaloc-runtime-pjrt") // GPU execution; runs on JDK 25 only
}

To build from source, run ./gradlew publishToMavenLocal -x test in a checkout and put mavenLocal() first in both repository blocks.

Walkthrough, plugin options, configuration and troubleshooting: GETTING_STARTED.md.

• Alpha. Names and signatures can change between alphas.
• f64 on Metal and Cerebras is f32. Apple GPUs and the WSE have no double precision; f64 programs run and return doubles with f32 accuracy.
• Cerebras is a work in progress. It runs on the SDK's simulator only, slowly, and three test classes still fail on programs too big for the current lowering.
• The CPU interpreter is for correctness, not speed. Performance numbers are for the GPU path.
• Tlaloc does not generate GPU code. XLA and IREE do. KPTX, a PTX kernel DSL, can replace a recognized op, but no kernel is on by default: its paged attention wins at large shapes and loses at small ones (KPTX_PAGED_PERF.md).
• One Kotlin release per Tlaloc version, because the plugin uses compiler internals.
• No Python API. Other languages use the StableHLO artifact.
• Not written yet or never run: TPU, multi-host training, more than one GPU per Triton model, preempting a live sequence, image input for Muse Glimmer.
• Symja (LGPL-3.0), used to differentiate loops whose trip count is not a compile-time constant, is optional and not in your dependency graph unless you add it.

Standalone projects under examples/. Each has its own Gradle build and skips by name when its hardware is missing.

Design documents are in docs/.

  your Kotlin            K2 plugin              :ir                 backends
  ───────────            ─────────              ───                 ────────
  grad { f }   ──────►   reverse transform ──► DXIR ──► StableHLO ──► PJRT  (CUDA, TPU*)
  typed tensors          on the lambda's IR     │       + Shardy  ──► IREE  (CPU, CUDA)
  named axes             (also: Kotlin source)  │                  ──► host interpreter
                                                └── recognize + coarsen (attention,
                                                    RMSNorm, RoPE, SwiGLU, …)

Loops are coarsened before differentiation, following Shen et al., Coarsening Optimization for Differentiable Programming (OOPSLA 2021).

./gradlew test                    # the suite
./gradlew test --rerun-tasks      # re-run everything
bash scripts/onboarding-smoke.sh  # publish, then build the README's install blocks
bash triton/verify.sh             # Triton backend, end to end (GPU + Docker)

Tests that need a GPU, Docker or a model download skip by name without it. Contributions: CONTRIBUTING.md. Security reports: SECURITY.md.

φ-calculus coarsening follows Shen, Zhang, Dea et al., Coarsening Optimization for Differentiable Programming, Proc. ACM Program. Lang. 5, OOPSLA, Article 130 (2021) — DOI 10.1145/3485507. The op surface and model layer take their parity target from facebookresearch/diffkt (MIT); the readable-derivative idea comes from google/tangent. Lowering targets OpenXLA StableHLO, Shardy and PJRT. Orchestration vendors Netflix Maestro.

Copyright 2026 Pedro N. Rodriguez. Licensed under Apache-2.0.

third-party/maestro/ vendors Netflix Maestro (Apache-2.0); NOTICE credits it. Tlaloc does not publish it. The paper text in docs/papers/ is CC-BY-4.0; see its README.

Symja (org.matheclipse:matheclipse-core), the optional CAS, is LGPL-3.0 by its published POM, while its upstream repository's license.txt is GPL-3.0; upstream's position is that the published Maven modules are LGPL. Tlaloc links it across the SymbolicEngine interface and never modifies or redistributes it. It is compileOnly, so it is not in your dependency graph unless you add it.