Kauldron ML Research Library

The core goal of konfig is to remove the indirection between your code and its configuration. Instead of creating custom configuration classes (like @dataclasses.dataclass) that duplicate parameters from your actual logic, you write regular Python code.

How it works: konfig treats a configuration as a nested tree of dict objects. There is a one-to-one mapping between nested Python function calls and their dictionary representations. When you wrap imports in konfig.imports(), calling a function doesn't execute it immediately; instead, it returns a konfig.ConfigDict that captures the function's __qualname__ and its arguments.

Key characteristics of konfig.ConfigDict:

• Serializable: You can convert it to a JSON-compatible dictionary using .to_json().
• Mutable: Unlike the original Python objects, ConfigDict objects can be modified (e.g., cfg.learning_rate = 1e-4).
• No Control Flow: Because they are dict-like, you can define nested calls like f(g()), but you cannot perform operations like f() + g() because that would attempt to add two dictionaries.

To turn a configuration dictionary back into a real Python object, use konfig.resolve(cfg).

# Instead of a real object, this returns a ConfigDict
with konfig.imports():
    import optax

cfg = optax.adam(learning_rate=1e-5)

# To get the real object back:
optimizer = konfig.resolve(cfg)

Use KTyping array types to annotate arrays for runtime checking of the array class, dtype, and shape. Annotations follow the format Type["shape_string"] (e.g., Float["b d"]).

Available Types

• General Purpose: Array (any array), Num (any numerical type), Float, Int, Bool, Complex.
• Specific Dtypes: UInt8, Float32, BFloat16, etc.
• Scalars: Scalar, ScalarInt, ScalarFloat, ScalarBool (supports 0-D arrays and Python scalars).
• Specialized: PRNGKey (JAX random keys).
• Framework Specific:
  • Tf... prefix (e.g., TfFloat) for TensorFlow.
  • X... prefix (e.g., XFloat) to support JAX, Numpy, TensorFlow, and Torch simultaneously.

Dimension Naming

• Common single letters: b or *b (batch), n (items/tokens), h/w (height/width), c (channels/features), d (embedding), t (time).
• Descriptive names: Use num_tokens or embedding_dim for clarity.
• Wildcards: Use ... or *any for dimensions that do not matter. *any is preferred if multiple arguments share the same dimensions.
• Anonymous dimensions: Prefix with _ (e.g., _b) to check the value without remembering it for later use. ... is a shortcut for *_.

To ensure compatibility with arbitrary batch sizes, follow these best practices:

1. Use Wildcard Batch Dimensions: Annotate with "*b ..." instead of a fixed dimension "b ...". This allows the code to handle any number of leading batch dimensions.
2. Use Negative Indexing: When indexing, always index from the right (e.g., img.mean((-3, -2, -1))) to remain agnostic to the number of batch dimensions.
3. Right-pad for Broadcasting: To allow implicit broadcasting between related arrays (like logits and labels), right-pad shapes with 1s. For example, if logits are Float["*b d"], labels should be Int["*b 1"] rather than Int["*b"].

# Good: Supports arbitrary batch shapes and uses negative indexing
def center(img: Float["*b h w c"]) -> Float["*b h w c":
  return img - img.mean((-3, -2, -1))

# Good: Right-padding for broadcasting
def ce(logits: Float["*b d"], labels: Int["*b 1"], mask: Bool["*#b 1"]) -> FloatScalar:
  ...

Directories are computed dynamically based on the root_dir provided in the kxm.Experiment:

• xp_dir: Experiment directory (default: {root_dir}/{xid}/)
• wu_dir: Work-unit directory (default: {root_dir}/{xid}/{wid}/)

Using the Work-unit Directory Proxy

To pass the work-unit directory as an argument to your binary, use kxm.WU_DIR_PROXY. This is a proxy string that resolves to the actual path once the experiment ID (xid) and work-unit ID (wid) are known.

Customizing Directory Names

Use kxm.SubdirFormat to customize the xp_dirname and wu_dirname. Supported placeholders include:

• {cell}: Cell from the main job
• {name} / {title}: Experiment title
• {user} / {author}: Author who launched the experiment
• {xid}: Experiment id
• {wid}: Work-unit id (padded)
• {unpadded_wid}: Work-unit id (unpadded)
• {sweep_kwargs}: String representation of the sweep (e.g. batch_size=32,seed=0)
• {separator_if_sweep}: A - token present only if there is a sweep

Example:

xp = kxm.Experiment(
    subdir_format=kxm.SubdirFormat(
        xp_dirname='{xid}-{name}',
        wu_dirname='{wid}-{sweep_kwargs}',
    ),
)

xp = kxm.Experiment(
    subdir_format=kxm.SubdirFormat(
        xp_dirname='{xid}-{name}',  # Default to `{xid}`
        wu_dirname='{wid}-{sweep_kwargs}',  # Default to `{wid}`
    ),
)

Kauldron uses a configuration system based on ml_collections.ConfigDict. Experiments are composed in a config.py file by choosing sub-modules. The root of the configuration is typically a kd.train.Trainer object.

To enable the config system, you must use the konfig.imports() context manager. Inside this context, standard library or third-party module calls (like optax.adam) do not return actual objects, but instead return ConfigDict builders. These builders create nested ConfigDict objects that represent the configuration of those modules.

To convert a configuration into the actual functional objects (e.g., turning a config for an optimizer into the optimizer itself), use konfig.resolve(cfg).

from kauldron import konfig
import optax

with konfig.imports():
  # This creates a ConfigDict object instead of a real optimizer
  cfg = optax.adam(learning_rate=0.003)

# Config can be mutated
cfg.learning_rate = 1e-4

# Resolve the actual object
optimizer = konfig.resolve(cfg)

To avoid creating NNX modules during config resolution, do not put NNX modules directly in Kauldron configs. Instead, use functools.partial in your config to pass an initialization function to a KdNnxModule.

Pattern:

1. In your config, assign a partial function to a field in your main KdNnxModule.
2. In the KdNnxModule.setup method, call that function to instantiate the sub-module.

This ensures module creation is delayed until the init phase.

# In your config:
with konfig.imports():
  from functools import partial
  from ... import my_module

def get_config():
  cfg = Trainer()
  # Use partial to delay instantiation
  cfg.model = my_module.MainModule(
    backbone_init=partial(my_module.NNXBackbone, arg1=value1)
  )
  return cfg

# In your my_module file:
@dataclasses.dataclass(kw_only=True)
class MyModule(knnx.KdNnxModule):
  backbone_init: tp.Callable

  def setup(self):
    self.backbone = self.backbone_init()

  def __call__(self, x):
    out = self.backbone(x)
    return out

Dimension assignments (e.g., b=32) are tracked using a stack of kt.ShapeScopes. Shape checks always refer to the active (topmost) scope on the stack, which can be accessed via kt.get_current_scope().

Managing Scopes

• Automatic Scopes: The @kt.typechecked decorator automatically opens a new scope when a function is called and closes it when the function returns. Arguments are checked against this scope, and internal checks (like kt.check_type) reuse and potentially modify this scope.
• Manual Scopes: You can manually open a scope using kt.typechecked as a context manager.
• Reusing Parent Scopes: To avoid opening a new scope and instead reuse or modify the parent scope, use @kt.typechecked(new_scope=False).

Scope Behavior with Dataclasses and TypedDicts

When a dataclass or TypedDict is decorated with @kt.typechecked, a new scope is opened specifically for checking its attributes during __init__ or when checking arguments annotated with that type. This means that if a function takes two arguments of the same decorated type, they each get their own independent scope (e.g., net1.logits.shape might differ from net2.logits.shape).

@kt.typechecked  # < Opens a new scope when fn is called
def fn(a: Float["a"]) -> Int["a"]:
  b, c = something(a)
  # all checks inside reuse this scope (and modify it if they succeed).
  kt.check_type(b, Float["a b"])

  # isinstance checks against the active scope, but does not modify it.
  if isinstance(c, Int["a c"]):
    ...
    d = fn(a[:3])

  return e # The return type is checked and then the scope is closed