Google Research has released Kauldron, a framework for machine learning experiments built on JAX. The system utilizes a configuration architecture where settings are stored as plain dictionaries, allowing for complex structures to be serialized to JSON without requiring specific base classes or decorators.
The library introduces cfg.ref for managing configuration dependencies. This feature allows parameters to reference other values dynamically, ensuring that updates to core settings, such as the number of training steps, automatically propagate through dependent schedules like learning rate decay curves.
Kauldron uses the kontext module to decouple components. By using string-based key paths to access data within a nested context, metrics and models can interact without direct imports. This architecture simplifies the process of redirecting inputs by updating a single string reference.
The ktyping module provides runtime shape validation for JAX arrays using named axes. By annotating function signatures with axis names, the framework enforces consistency across arguments and return values. In case of a mismatch, the system identifies the specific axis name that caused the error.
The framework supports custom losses and metrics defined as frozen dataclasses. Metrics utilize an AutoState mechanism with sum_field entries to ensure accurate aggregation across batches and devices. This approach prevents errors that typically occur when averaging rates from batches of different sizes.
