Kauldron от Google Research: JAX-библиотека для конфигов, которые не ломаются, и тренера, который читается от начала до конца

Kauldron — это не очередной фреймворк, а три механизма, которые Google Research собрала в одну библиотеку. Первый — konfig: весь эксперимент превращается в дерево обычных словарей, которые сериализуются в JSON и обратно. Второй — kontext: компоненты не импортируют друг друга, а соединяются строковыми ключами — loss никогда не знает, какую модель он оценивает. Третий — runtime shape checker: именованные оси (например, Float['*b h w c']) проверяются на лету, и если размерность не совпала, он сообщает, какие имена были к чему привязаны.

В туториале автор устанавливает Kauldron 1.4.2, применяет патч совместимости для JAX 0.10.1 (иначе etils ломается на первом шаге), и проходит весь пайплайн: пишет кастомную loss-функцию и метрику, тренирует модель на синтетических данных без скачивания датасета и без ускорителя, мониторит внутренний слой модели без её редактирования. В конце — запуск пяти вариантов эксперимента, где каждый отличается одной строкой конфига, и чекпоинт с возобновлением обучения.

Kauldron позиционируется как библиотека для research velocity: конфиги — это данные, а не код, компоненты не зависят друг от друга, а тренер читается от начала до конца. Для тех, кто устал от многоэтажных конфигов в Flax и Optax, это может быть глотком свежего воздуха. Но пока библиотека требует патча совместимости и работает только на CPU в демо.

Инструмент доступен на GitHub под лицензией Apache 2.0. Для JAX-разработчиков, которые хотят писать эксперименты без boilerplate и с полной прозрачностью, Kauldron — кандидат на замену текущим стекам.