Google's Kauldron lets you swap any dataset, model, or optimizer into one training pipeline

Kauldron: Modular, scalable library to train ML models

Kauldron is a Google Research library for training ML models, built on JAX and Flax, that prioritizes research velocity and modularity. Every component is self-contained and usable outside the framework, so you can plug in any dataset (TFDS, Grain, SeqIO, or your own), any Flax model, and any optimizer, then override anything from model architecture to inner layer parameters. It ships with ready-to-run example configs, a Colab-first workflow, and built-in profiling and debugging tools. Note: it is not an officially supported Google product.