Google Research представила Kauldron — фреймворк для машинного обучения на базе JAX. Система использует конфигурации в виде обычных словарей, что позволяет сериализовать сложные структуры в JSON без необходимости внедрения специальных базовых классов или декораторов в сторонние библиотеки.
Для управления зависимостями параметров в Kauldron предусмотрен механизм cfg.ref. Он позволяет связывать значения между собой, благодаря чему изменения в ключевых настройках, например, в количестве шагов обучения, автоматически пересчитывают зависимые графики, такие как кривые затухания скорости обучения.
Модуль kontext обеспечивает слабую связанность компонентов. Взаимодействие между моделями и метриками происходит через строковые пути доступа к данным в общем контексте, что исключает необходимость прямого импорта. Перенаправление данных осуществляется путем изменения одной строковой ссылки.
Система ktyping выполняет проверку размерностей массивов JAX во время выполнения с помощью именованных осей. Аннотации в сигнатурах функций позволяют отслеживать соответствие размерностей между аргументами и возвращаемыми значениями. При ошибке фреймворк указывает конкретное имя оси, вызвавшей конфликт.
Пользовательские функции потерь и метрики в Kauldron оформляются как неизменяемые дата-классы. Метрики используют механизм AutoState с полями суммирования, что гарантирует точность агрегации данных при работе с батчами разного размера и распределенных вычислениях, предотвращая ошибки при усреднении результатов.
