A Coding Guide to Google Research’s Kauldron: Configs That Are Plain Data, Components Wired by String, and a JAX Trainer You Can Read End to End

This tutorial is less about “here’s another training stack” and more about what Kauldron actually changes in day-to-day work. The pitch is simple: research velocity and modularity, but the notebook keeps testing that promise against real code instead of slogans. It installs Kauldron, patches one current compatibility snag, and then spends most of its time on the plumbing that makes the library feel different from a typical Flax plus Optax setup.

The first piece is konfig. Inside a konfig.imports() block, a call like optax.adam(...) doesn’t build an optimizer yet; it builds a plain config object that can be mutated, serialized to JSON, and later resolved into the real thing. That matters because the config is just data. The same pattern works for a more complex optax.chain too, and the tutorial makes a point of stressing that optax itself doesn’t need to know anything about Kauldron to participate.

Then comes cfg.ref, which is the part that keeps sweeps from lying to you. A warmup-cosine schedule points its decay_steps at cfg.ref.num_train_steps, not a copied-in literal, and when num_train_steps changes from 1000 to 200 the schedule changes with it. That sounds small. It isn’t. This is how you avoid quietly training every variant in a sweep on a stale decay curve.

The wiring story is kontext. A context is just nested data, and key paths like batch.image or preds.aux[0].pos reach into it. Metrics and other components declare their inputs with kontext.Key fields, then resolve_from_keyed_obj pulls the right tensors out as keyword arguments. The metric never imports the model, the model never imports the metric, and changing what a metric reads is often just changing a string.

Kauldron’s runtime shape checker gets the same treatment. The tutorial uses named axes in type annotations, binds them at runtime, and shows the failure mode when the shapes disagree: the error reports the axis name that was already bound, not just two anonymous tuples. After that, the notebook builds a custom log-cosh loss and a custom metric, runs a Trainer on synthetic in-memory data, watches an inner layer without editing the model, runs a five-variant sweep where each run differs by one config line, and finally lets a training job checkpoint itself and resume from where it stopped.