### Initialize Trainer and Launcher

Source: https://github.com/google-research/kauldron/blob/main/docs/konfig_philosophy.md

Examples of initializing top-level abstractions for training and experiment launching.

```python
trainer = kd.train.Trainer(
    train_ds=ds,
    model=model,
    optimizer=optimizer,
    ...,
)
trainer.train()
```

```python
xp = kxm.Experiment(
    name='My experiment',
    jobs={
        'train': kxm.Job(
            target='//path/to/my:target',
            platform='df=2x2',
        ),
    },
)
xp.launch()
```

--------------------------------

### Standard Optimizer Initialization

Source: https://github.com/google-research/kauldron/blob/main/docs/konfig_philosophy.md

A basic example of initializing an optimizer using optax.

```python
optimizer = optax.adam(learning_rate=1e-5)
```

--------------------------------

### Migrate from jaxtyping to ktyping

Source: https://github.com/google-research/kauldron/blob/main/kauldron/ktyping/README.md

Example showing the transition from jaxtyping decorators and annotations to ktyping's simplified syntax.

```python
# Before:
from jaxtyping import Array, Float, Int, jaxtyped  # pylint: disable=g-multiple-import, g-importing-member
import typeguard

@jaxtyped(typechecker=typeguard.typechecked)
def fn(x: Float[Array, "b d"], y: Integer[Array, "b"]) -> Float[Array, "b d"]:
    ...

# After:
import kauldron.ktyping as kt
from kauldron.ktyping import Float, Int, typechecked  # pylint: disable=g-multiple-import, g-importing-member

@kt.typechecked
def fn(x: Float["b d"], y: Int["b"]) -> Float["b d"]:
    ...
```

--------------------------------

### Launch Experiments with XManager

Source: https://github.com/google-research/kauldron/blob/main/CHANGELOG.md

Example of launching a Kauldron experiment using XManager. This command allows for deep configuration, sweeping parameters, and specifying hardware platforms.

```sh
xmanager launch third_party/py/kauldron/xm/launch.py -- \
      --cfg=third_party/py/kauldron/examples/mnist_autoencoder.py \
      --cfg.train_ds.batch_size=32 \
      --xp.sweep \
      --xp.platform=a100 \
      --xp.debug.catch_post_mortem
```

--------------------------------

### End-to-end konfig example

Source: https://github.com/google-research/kauldron/blob/main/kauldron/konfig/docs/demo.ipynb

Demonstrates the basic workflow of konfig: importing modules, creating a configuration, mutating it, and resolving it into a Python object.

```python
# Step 1: Import your modules you want to be configurable (can be any module)
with konfig.imports():
  from flax import linen as nn

# Step 2: Create your configuration. **Nothing is executed yet!!**
model_cfg = nn.Dense(features=32)  # This create a `ConfigDict` !!!

# Step 3: The `ConfigDict` object can optionally be mutated
model_cfg.use_bias = True

# Step 4: Resolve the `ConfigDict` -> Python object
# This is where the actual `nn.Dense` object is created
model = konfig.resolve(model_cfg)
```

--------------------------------

### Konfig Optimizer Configuration

Source: https://github.com/google-research/kauldron/blob/main/docs/konfig_philosophy.md

Examples of configuring optimizers directly using Konfig syntax.

```python
cfg = optax.adam(learning_rate=1e-5, b2=0.99)
```

```python
cfg = optax.chain(
    optax.scale_by_adadelta(),
    optax.add_decayed_weights(weight_decay=0.0),
    optax.scale_by_learning_rate(learning_rate=0.003),
)
```

--------------------------------

### Configure PyGrain Data Source

Source: https://github.com/google-research/kauldron/blob/main/docs/data.md

Example of configuring a PyGrain data source using kd.data.py.DataSource. This setup includes specifying the data source, enabling shuffling, setting batch size, and applying transformations like element filtering and value range normalization.

```python
cfg.train_ds = kd.data.py.DataSource(
    data_source=tfds.data_source("mnist", split='train'),
    shuffle=True,
    batch_size=256,
    transforms=[
        kd.data.Elements(keep=["image"]),
        kd.data.ValueRange(key="image", vrange=(0, 1)),
    ],
)
```

--------------------------------

### Minimal tf.data Pipeline Example

Source: https://github.com/google-research/kauldron/blob/main/kauldron/data/tf/README.md

Use this for a basic tf.data pipeline configuration with TFDS. Ensure the dataset is in ArrayRecord format.

```python
cfg.train_ds = kd.data.tf.Tfds(
    # TFDS parameters
    name='mnist',
    split='train',

    # `kd.data.tf.TFDataPipeline` optional parameters (common to all objects)
    batch_size=32,
    transforms=[
        kd.data.Elements(keep=["image"]),
    ],
)
```

--------------------------------

### Dim Assignments Example

Source: https://github.com/google-research/kauldron/blob/main/kauldron/ktyping/README.md

Shows how dimension assignments are inferred and displayed when type checking.

```none
Dim Assignments:
 - {*b: (32, 128), d: 1000}
```

--------------------------------

### Adding Transformations to a Custom Source

Source: https://github.com/google-research/kauldron/blob/main/kauldron/data/tf/README.md

Example of how to add transformations to a custom dataset configuration. Transformations should be passed via the `transforms` argument.

```python
cfg.train_ds = MyDataset(
    transforms=[
        MyDatasetTransform(),
        kd.data.tf.Resize(key='image', height=32, width=32),
    ],
)
```

--------------------------------

### Configure Hierarchical NNX Modules with Partial

Source: https://github.com/google-research/kauldron/blob/main/docs/nnx.md

Example of using `functools.partial` in Kauldron configurations to delay the instantiation of NNX modules. This prevents NNX modules from being created during config resolution.

```python
with konfig.imports():
  from functools import partial
  from ... import my_module

def get_config():
  cfg = Trainer()
  cfg.model = my_module.MainModule(
  backbone_init=partial(my_module.NNXBackbone, arg1=value1))
  ...
  return cfg
```

```python
from flax import nnx
from kauldron.contrib import knnx
import typing as tp

class NNXBackbone(nnx.Module):
  def __init__(self, arg1):
    self.arg1 = value1

# main kauldron module
@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)
    ...

```

--------------------------------

### Traditional Config Class Pattern

Source: https://github.com/google-research/kauldron/blob/main/docs/konfig_philosophy.md

An example of the boilerplate-heavy approach to configuration that Konfig aims to replace.

```python
@dataclasses.dataclass
class AdamOptimizerConfig:
  learning_rate: float = 1e-5

  def make(self):
    return optax.adam(learning_rate=self.learning_rate)
```

--------------------------------

### TypeCheckError Example

Source: https://github.com/google-research/kauldron/blob/main/kauldron/ktyping/README.md

Illustrates a TypeCheckError from jaxtyping, showing the mismatch between actual and expected types for a function parameter.

```none
TypeCheckError: jax.jaxlib._jax.ArrayImpl of "argument 'mask'" (jax.jaxlib._jax.ArrayImpl)
did not match any element in the union:
  jaxtyping.Bool[Array, '*b 1']: is not an instance of jaxtyping.Bool[Array, '*b 1']
  NoneType: is not an instance of NoneType

During handling of the above exception, another exception occurred:

[...] 

TypeCheckError:
The problem arose whilst typechecking parameter 'mask'.
Actual value: bool[32,1](jax)
Expected type: jaxtyping.Bool[Array, '*b 1'] | None.

The above exception was the direct cause of the following exception:

[...] 

TypeCheckError: Type-check error whilst checking the parameters of __main__.cross_entropy_loss.
The problem arose whilst typechecking parameter 'mask'.
Actual value: bool[32,1](jax)
Expected type: jaxtyping.Bool[Array, '*b 1'] | None.
----------------------
Called with parameters: {'preds': f32[32,128,1000](jax), 'targets': i32[32,128,1](jax),
 'mask': bool[32,1](jax)}
Parameter annotations: (preds: Float[Array, '*b d'], targets: Integer[Array, '*b 1'],
 mask: Bool[Array, '*b 1'] | None = None) -> Any.
The current values for each jaxtyping axis annotation are as follows.
d=1000
b=(32, 128)
```

--------------------------------

### Multiple Dim Assignment Candidates Example

Source: https://github.com/google-research/kauldron/blob/main/kauldron/ktyping/README.md

Demonstrates a scenario where multiple valid dimension assignments can be inferred, leading to ambiguity.

```python
@kt.typechecked
def foo(x: Float["b n"] | Float["b m"], y: Float["b m"]):
  ...

foo(np.zeros((8, 3)), np.zeros((7, 3)))
```

```none
Multiple Dim Assignment Candidates:
 - {b: 8, n: 3}
 - {b: 8, m: 3}
```

--------------------------------

### Temporarily Modify Ktyping Config

Source: https://github.com/google-research/kauldron/blob/main/kauldron/ktyping/README.md

Example of using a context manager to temporarily disable type checking.

```python
import kauldron.ktyping as kt

with kt.Config(typechecking_enabled=False):
  ...
```

--------------------------------

### Define Custom Kauldron NNX Module

Source: https://github.com/google-research/kauldron/blob/main/docs/nnx.md

Inherit from `kd.contrib.knnx.KdNnxModule` and define NNX module components in the `setup` method. This allows NNX modules to be used within Kauldron's configuration system by delaying instantiation.

```python
from kauldron import kd
import dataclasses

@dataclasses.dataclass(kw_only=True)
class MyKdNnxModule(kd.contrib.knnx.KdNnxModule):
  input_dim: int = 3
  hdim: int = 10

  image: kontext.Key = "batch.image"

  def setup(self, rngs: nnx.Rngs=nnx.Rngs(0)):
    self.lin = nnx.Linear(self.input_dim, self.hdim, rngs=rngs)

  def __call__(self, image):
    return self.lin(image)
```

--------------------------------

### KTyping f-string Expression Example

Source: https://github.com/google-research/kauldron/blob/main/kauldron/ktyping/README.md

Illustrates the use of f-string expressions within shape specifications for dynamic dimension sizes. The 'features' argument is used to define the output shape.

```python
@kt.typechecked
def embed(x: Float["b n d"], features: int) -> Float["b n {features}"]:
  ...
```

--------------------------------

### Defining complex optimizers

Source: https://github.com/google-research/kauldron/blob/main/docs/konfig_philosophy.md

Examples of simple and complex optimizer configurations, emphasizing the need to move complex abstractions into Python code rather than config files.

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

```python
cfg = optax.chain(
    optax.scale_by_adadelta(),
    optax.add_decayed_weights(weight_decay=0.0),
    optax.scale_by_learning_rate(learning_rate=0.003),
)
```

```python
cfg = my_project.my_complex_optimizer(learning_rate=0.003)
```

--------------------------------

### Get key paths from a KeyedObject using kontext.get_keypaths

Source: https://github.com/google-research/kauldron/blob/main/kauldron/kontext/README.md

Use `kontext.get_keypaths` to retrieve the path values defined by `Key` annotations in a `KeyedObject`.

```python
a = A(x='a[0].b.inner', y=None)

assert kontext.get_keypaths(a) == {
      'x': 'a[0].b.inner',
      'y': None,
  }
```

--------------------------------

### Relaunch Experiment with PartialKauldronLoader

Source: https://github.com/google-research/kauldron/blob/main/docs/checkpoint.md

Initialize a new experiment from a previous state by mapping core components like step, params, and optimizer state.

```python
cfg.init_transform = kd.ckpts.PartialKauldronLoader(
    workdir=/path/to/old/workdir/,
    new_to_old={  # Mapping params
        'step': 'step',
        'params': 'params',
        'collections': 'collections',
        'opt_state': 'opt_state',
    },
)
```

--------------------------------

### Define a nested tree structure

Source: https://github.com/google-research/kauldron/blob/main/kauldron/kontext/README.md

This is an example of a nested dictionary structure that can be manipulated by Kontext.

```python
tree = {
    'a': [{'b': {'inner': 123}, 'c': 456}, {'e': 789}]
}
```

--------------------------------

### Initialize and Use PRNGKey

Source: https://github.com/google-research/kauldron/blob/main/kauldron/random/README.md

Demonstrates initializing a PRNGKey and splitting it using the OO API. Supports folding in strings.

```python
key = kd.random.PRNGKey(0)

key0, key1 = key.split()  # ObjecO API

key1 = key1.fold_in('param')  # Fold-in supports `str`

# Those 2 lines are equivalent:
x = key0.uniform()
x = jax.random.uniform(key)  # Jax API still works
```

--------------------------------

### Handle KTypeCheckError

Source: https://github.com/google-research/kauldron/blob/main/kauldron/ktyping/README.md

Example of a function signature that triggers a KTypeCheckError when arguments do not match the specified shape.

```python
@kt.typechecked
def cross_entropy_loss(
    preds: Float["*b d"], targets: Int["*b 1"], mask: Bool["*b 1"] | None = None
) -> ScalarFloat:
  ...
```

```none
[...]

KTypeCheckError: argument mask = np.bool_[32 1] has shape (32, 1) which is not
shape-compatible with '*b 1' (required by Bool['*b 1'] | None).

Origin: function 'cross_entropy_loss' at /tmp/ipython-input-3-3053638833.py:1

Arguments:
  preds: Float['*b d'] = np.f64[32 128 1000]
  targets: Int['*b 1'] = np.i64[32 128 1]
> mask: Bool['*b 1'] | None = np.bool_[32 1]

Return: ScalarFloat

Dim Assignments:
 - {*b: (32, 128), d: 1000}
```

--------------------------------

### Instantiate and Use Custom NNX Module

Source: https://github.com/google-research/kauldron/blob/main/docs/nnx.md

Demonstrates how to instantiate and use a custom Kauldron NNX module, both as a direct NNX module and as a Flax Linen module.

```python
mod = MyModule()
mod.setup()
out = mod(jnp.ones(1, 32, 32, 3))
```

```python
mod = MyModule()
vars = mod.init({'params':jax.random.key(0)}, jnp.ones(1, 32, 32, 3))
out = mod.apply(vars, rngs={})
```

--------------------------------

### Minimal XManager experiment using Direct API

Source: https://github.com/google-research/kauldron/blob/main/kauldron/xm/README.md

Define and launch a minimal experiment using the direct API. Jobs can be configured with target, platform, and arguments.

```python
from kauldron import kxm

xp = kxm.Experiment(
    jobs={
        'train': kxm.Job(
            target='//path/to/my:trainer',
            platform='jf=2x2',
            args={
                'workdir': kxm.WU_DIR_PROXY,
            },
        )
    },
    root_dir='/path/to/home/{author}/experiments/',
)
xp.launch()
```

--------------------------------

### Launch a single binary via CLI using existing config

Source: https://github.com/google-research/kauldron/blob/main/kauldron/xm/README.md

Use this command to launch a single binary using a predefined configuration. Specify the target binary directly.

```sh
xmanager launch kauldron/xm/launch.py -- \
  --xp=kauldron/xm/configs/single.py \
  --xp.target=//path/to/my:binary
```

--------------------------------

### Import Kauldron Metrics, Losses, and Summaries

Source: https://github.com/google-research/kauldron/blob/main/docs/metrics.md

Import necessary modules for using metrics, losses, and summaries in standalone applications.

```python
from kauldron import metrics
from kauldron import losses
from kauldron import summaries
```

--------------------------------

### Propagate Configuration Flags

Source: https://github.com/google-research/kauldron/blob/main/kauldron/xm/README.md

Use kxm.CFG_FLAG_VALUES to pass configuration flags from the launch script to the binary.

```python
kxm.Job(
    args={
        'cfg': kxm.CFG_FLAG_VALUES,
    },
)
```

```sh
xmanager launch kauldron/xm/launch.py -- \
  --cfg=path/to/my_config.py \
  --cfg.train_ds.batch_size=32 \
  --xp=path/to/xm/launch_config.py \
  --xp.sweep=True \
  --xp.platform=jf=2x2
```

--------------------------------

### Restore Trainer from Workdir

Source: https://github.com/google-research/kauldron/blob/main/docs/checkpoint.md

Initialize a trainer by resolving a configuration file stored in a previous experiment's workdir.

```python
config_path = pathlib.Path(workdir) / 'config.json'
trainer = kd.konfig.resolve(json.loads(config_path.read_text()))
```

--------------------------------

### Manage File Dependencies

Source: https://github.com/google-research/kauldron/blob/main/kauldron/xm/README.md

Package and resolve file paths dynamically using proxy strings and the files dictionary.

```python
kxm.Job(
    args={
        # Value dynamically resolved later
        'gin_config': kxm.file_path('my_file.gin'),
    },
    # Files to package with the MPM
    files={
        'my_file.gin': '//path/to/file.gin'
    },
)
```

--------------------------------

### Import and Configure Datasets in Konfig

Source: https://github.com/google-research/kauldron/blob/main/kauldron/konfig/README.md

Illustrates how to structure Konfig projects by separating config files from implementation code. Configuration files should reside in the `configs/` directory and be imported within the `konfig.imports()` context.

```python
from my_project.configs import base_config
from my_project.configs import ds_config

with konfig.imports():
  from my_project import my_ds


def get_config():
  cfg = base_config.get_config()
  cfg.train_ds = ds_config.make_ds(split='train')  # This returns a `ConfigDict`
  cfg.eval_ds = my_ds.MyDataset()  # This IS a `ConfigDict`
```

--------------------------------

### Start Eval-Only Job

Source: https://github.com/google-research/kauldron/blob/main/docs/eval.md

Use `kd.train.Trainer.eval_only()` to run evaluation on a trainer from a previous Kauldron experiment. This method only works within `konfig`.

```python
def config():
  return kd.train.Trainer.eval_only(
      evals = {
          'my_eval': kd.evals.Evaluator(
              run=kd.evals.StandaloneLastCheckpoint(),
              ...,
          ),
      }
  )
```

--------------------------------

### Partial Loading with init_transform

Source: https://github.com/google-research/kauldron/blob/main/docs/checkpoint.md

Use PartialKauldronLoader to map and load specific subtrees from a pretrained checkpoint during initialization.

```python
cfg.init_transform = kd.ckpts.PartialKauldronLoader(
    workdir=/path/to/old/workdir/,
    new_to_old={  # Mapping params
        # The new_to_old dict determines which weights are loaded from the
        # target checkpoint, and can also be used to rename subtrees when
        # loading params from a different pretrained model.
        # '<new_path>':            '<source_path>'
        'params.decoder.layers_0': 'params.encoder',
    },
)

trainer = konfig.resolve(cfg)

# When initializing the weights, the `init_transform` is applied
init_state = trainer.init_state()

# `init_state.params['decoder']['layers_0']` now contains the previous encoder
# weights
```

--------------------------------

### Basic Usage of KTyping

Source: https://github.com/google-research/kauldron/blob/main/kauldron/ktyping/README.md

Demonstrates basic usage of KTyping with type-checked function arguments and return types. Imports required libraries and defines a function with annotated array types.

```python
import jax.numpy as jnp
import jax.nn as nn
import kauldron.ktyping as kt
from kauldron.ktyping import Float

@kt.typechecked
def simple_attention(
    query: Float["*b n d"],
    key: Float["*b m d"],
    value: Float["*b m dv"],
    mask: Bool["*b n m"]
) -> Float["*b n dv"]:
  attn_weights = jnp.einsum("...nd,...md->...nm", query, key)
  attn_weights = nn.softmax(attn_weights, axis=-1)
  kt.check_type(attn_weights, Float["*b n m"])
  return jnp.einsum("...nm,...mv->...nv", attn_weights*mask, value)
```

--------------------------------

### Compute Metric via State (Single Step)

Source: https://github.com/google-research/kauldron/blob/main/docs/metrics.md

An alternative way to compute a metric for a single step by first getting its state and then computing the final value.

```python
accuracy = metric.get_state(logits=logits, labels=labels).compute()
```

--------------------------------

### Launch experiment with CLI parameter overwrites

Source: https://github.com/google-research/kauldron/blob/main/kauldron/xm/README.md

Launch a configuration-based experiment and overwrite specific parameters like job platform, cell, and experiment notes via the CLI.

```sh
xmanager launch kauldron/xm/launch.py -- \
  --xp=path/to/my/xp_config.py \
  --xp.jobs.train.platform=df=4x4 \
  --xp.jobs.eval.platform=cpu \
  --xp.cell=jn \
  --xp.note="My experiment description"
```

--------------------------------

### Define Experiment Work Directories

Source: https://github.com/google-research/kauldron/blob/main/kauldron/xm/README.md

Use kxm.WU_DIR_PROXY to dynamically resolve the work-unit directory path at runtime.

```python
xp = kxm.Experiment(
    # The root dir is defined at the `Experiment` level
    root_dir='/path/to/home/{author}/kd/'
    jobs={
        'train': kxm.Job(
            # A proxy string is used, that will later be resolved as
            # --workdir="/path/to/home/{author}/kd/{xid}/{wid}/"
            args={'workdir': kxm.WU_DIR_PROXY},
        )
    },
)
```

--------------------------------

### Sweep ConfigArgs Values

Source: https://github.com/google-research/kauldron/blob/main/kauldron/konfig/README.md

Demonstrates how to sweep through different values of ConfigArgs for hyperparameter tuning. Each yielded dictionary represents a configuration to be tested.

```python
def sweep_cfgarg():
  for dataset_name in ["mnist", "imagenet"]:
    yield {"__args__.dataset_name": dataset_name}
```

--------------------------------

### Custom Array Type Definition

Source: https://github.com/google-research/kauldron/blob/main/kauldron/ktyping/README.md

Shows how to define custom array types using ArrayTypeMeta, specifying allowed array classes and dtypes. This example creates a type that accepts JAX and TensorFlow arrays with float32 or boolean dtypes.

```python
from kauldron.ktyping import array_type_meta as atm, dtypes

VeryUsefulArray = atm.ArrayTypeMeta(
  "VeryUsefulArray",
  array_types=(atm.JaxArray, atm.TfArray), # accept jax and tf (but not numpy)
  dtype=dtypes.float32 | dtypes.bool)  # either float32 or bool

def foo(x: VeryUsefulArray["a b c"]):
  ...
```

--------------------------------

### Launch a Kauldron experiment via CLI

Source: https://github.com/google-research/kauldron/blob/main/kauldron/xm/README.md

Use this command to launch a Kauldron experiment with specific configurations. It allows overriding parameters like batch size and platform.

```sh
xmanager launch kauldron/xm/launch.py -- \
  --cfg=kauldron/examples/mnist_autoencoder.py \
  --cfg.train_ds.batch_size=32 \
  --xp.sweep=True \
  --xp.platform=jf=2x2
```

--------------------------------

### Use Config References

Source: https://github.com/google-research/kauldron/blob/main/docs/intro.md

Enables synchronized updates across multiple configuration fields by using the .ref. attribute.

```python
cfg.num_train_steps = 10_000

cfg.schedules = {
    "learning_rate": optax.warmup_cosine_decay_schedule(
        init_value=0.0,
        peak_value=0.001,
        decay_steps=cfg.ref.num_train_steps,  # <<< `.ref.` used here !!!!!
    )
}
```

--------------------------------

### General Kauldron CLI Usage

Source: https://github.com/google-research/kauldron/blob/main/kauldron/cli/README.md

This is the general syntax for using the Kauldron CLI. Specify the command, sub-command, and configuration file path. Overrides can be applied using '--cfg.override.key=value'.

```sh
kauldron <command> <sub_command> \
    --cfg=<path/to/config.py> [--cfg.override.key=value ...]
```

--------------------------------

### Create a Metric Instance

Source: https://github.com/google-research/kauldron/blob/main/docs/metrics.md

Instantiate a metric object for standalone use. Metrics are stateless objects.

```python
metric1 = metrics.Accuracy()
```

--------------------------------

### Importing and Using KTyping Annotations

Source: https://github.com/google-research/kauldron/blob/main/kauldron/ktyping/STYLE.md

Recommended approach for importing ktyping types directly to improve annotation readability compared to prefixing every type with the module alias.

```python
import kauldron.ktyping as kt
from kauldron.ktyping import Float, UInt8  # pylint: disable=g-multiple-import,g-importing-member

@kt.typechecked
def fn(x: Float["b n c"]) -> UInt8["n n"]:
  n = kt.dim["n"]
  return np.zeros((n, n), dtype=np.uint8)
```

```python
import kauldron.ktyping as kt

@kt.typechecked
def fn(x: kt.Float["b n c"]) -> kt.UInt8["n n"]:
  n = kt.dim["n"]
  return np.zeros((n, n), dtype=np.uint8)
```

--------------------------------

### Define experiment config using get_config function

Source: https://github.com/google-research/kauldron/blob/main/kauldron/xm/README.md

Wrap your kxm.Experiment definition within a get_config function for better re-usability and CLI parameter overwriting.

```python
from kauldron import konfig

with konfig.imports():
  from kauldron import kxm

def get_config() -> kxm.Experiment:
  return kxm.Experiment(
      ...
  )
```

--------------------------------

### Add Module-Specific Ktyping Config Override

Source: https://github.com/google-research/kauldron/blob/main/kauldron/ktyping/README.md

Demonstrates how to apply a configuration override for a specific module using a regex.

```python
import kauldron.ktyping as kt

kt.add_config_override(
  r"example\.module",
  kt.config.Config(typechecking_enabled=False)
)
```

--------------------------------

### Create a config using konfig.imports

Source: https://github.com/google-research/kauldron/blob/main/kauldron/konfig/docs/demo.ipynb

An alternative method to create a `ConfigDict` by importing modules within a `konfig.imports()` context manager.

```python
with konfig.imports():
  from flax import linen as nn

model_cfg = nn.Sequential(layers=[
    nn.Dense(features=32),
    nn.Dropout(rate=0.5),
]);
```

--------------------------------

### Using a tf.data Pipeline Iterator

Source: https://github.com/google-research/kauldron/blob/main/kauldron/data/tf/README.md

Demonstrates how to use a configured tf.data pipeline as a standalone iterator. Includes device placement for sharding.

```python
ds = kd.data.tf.Tfds(...)
ds = ds.device_put(kd.sharding.FIRST_DIM)

for ex in ds:
  ...
```

--------------------------------

### Define ConfigArgs for Modular Configurations

Source: https://github.com/google-research/kauldron/blob/main/kauldron/konfig/README.md

Use ConfigArgs to declare main configuration hyperparameters for modular configs. Arguments can be set via the command line using the `--cfg.__args__.arg_name = value` pattern.

```python
import dataclasses

_DATASET_DIMS = {'mnist': 10, 'imagenet': 1000}

@dataclasses.dataclass(kw_only=True, frozen=True)
class ConfigArgs:
  dataset_name: str = 'mnist'
  model_name = 'vit'

def get_config(args: ConfigArgs = ConfigArgs()):
  cfg = kd.train.Trainer()
  if args.model_name == 'vit':
    cfg.model = MyModel(dim=_DATASET_DIMS[args.dataset_name])
  else:
    ...
```

--------------------------------

### Create and Resolve Kauldron Trainer Config

Source: https://github.com/google-research/kauldron/blob/main/docs/intro.md

Use `kd.konfig.mock_modules()` to create a mutable ConfigDict that mimics the Trainer object for IDE benefits. Resolve the ConfigDict into an immutable Trainer object using `kd.konfig.resolve()`.

```python
from kauldron import kd

with kd.konfig.mock_modules():
  # Inside mock_modules(), modules become ConfigDict builders
  # For the IDE, the ConfigDict look like the real object, so we get
  # all auto-complete and type-checking benefits.
  cfg = kd.train.Trainer()
  cfg.workdir = '/tmp/'

trainer = kd.konfig.resolve(cfg)  # ConfigDict becomes Trainer

assert isinstance(cfg, kd.konfig.ConfigDict)
assert isinstance(trainer, kd.train.Trainer)
```

--------------------------------

### Implement multi-optimizer training

Source: https://github.com/google-research/kauldron/blob/main/docs/eval.md

Use multi_optimizer and MultiTrainStep to manage separate optimizers for different model components.

```python
trainer = kd.train.Trainer(
    ...,
    model=MyGan(
        generator=MyGenerator(),
        discriminator=MyDiscriminator(),
    ),
    # Define the loss for the generator and the discriminator
    losses={
        'discriminator': kd.losses.L2(...),
        'generator': kd.losses.L2(...),
    },
    optimizer=kd.contrib.train.multi_optimizer(
        # Using `kd.optim.partial_update`, you can mask out which weights
        # each of the optimizer will be applied too.
        discriminator=kd.optim.partial_update(
            optimizer=optax.adam(1e-4),
            mask=kd.optim.select('discriminator'),
        ),
        generator=kd.optim.partial_update(
            optimizer=optax.adam(1e-4),
            mask=kd.optim.select('generator'),
        ),
    ),
    # Using `kd.contrib.train.multi_optimizer` require to use `MultiTrainStep`
    trainstep=kd.contrib.train.MultiTrainStep(),
)
```

--------------------------------

### Configure Evaluator Run Behavior

Source: https://github.com/google-research/kauldron/blob/main/CHANGELOG.md

Demonstrates how to configure when an evaluator runs. Use `kd.evals.RunEvery` to run periodically with the training job, or `kd.evals.RunXM` to launch in a separate evaluation job.

```python
cfg.evals = {
    'eval_train': kd.evals.Evaluator(
        run=kd.evals.RunEvery(100),  # Run along `train`
    ),
    'eval_eval': kd.evals.Evaluator(
        run=kd.evals.RunXM(),  # Run in a separate `eval` job.
    ),
}
```

--------------------------------

### Mid-Level Training Loop

Source: https://github.com/google-research/kauldron/blob/main/docs/eval.md

For direct control over the training loop, initialize the state and iterate through batches using `trainer.trainstep.step`. Use `.device_put()` to place batches on devices.

```python
state = trainer.init_state()

for batch in trainer.train_ds.device_put(trainer.sharding.batch):
  state, aux = trainer.trainstep.step(state, batch)
```

--------------------------------

### Launch an arbitrary Python script via CLI

Source: https://github.com/google-research/kauldron/blob/main/kauldron/xm/README.md

Use this command to launch any custom Python script. Specify the target binary and desired platform.

```sh
xmanager launch kauldron/xm/launch.py -- \
  --xp=kauldron/xm/configs/single.py \
  --xp.target=//path/to/my:binary \
  --xp.platform=df=2x2
```

--------------------------------

### Configure Sharding Strategy

Source: https://github.com/google-research/kauldron/blob/main/docs/sharding.md

Define sharding for model parameters and optimizer states using the ShardingStrategy configuration.

```python
cfg.sharding = kd.sharding.ShardingStrategy(
    params={
        'encoder': kd.sharding.REPLICATED,
        'decoder': my_project.my_sharding_strategy,
    },
    opt_state=None,  # Let jax auto-infer the sharding
)
```

--------------------------------

### Implement post-resolve hook

Source: https://github.com/google-research/kauldron/blob/main/kauldron/konfig/docs/demo.ipynb

Uses the __post_konfig_resolve__ protocol to track the configuration object after resolution.

```python
class A:

  def __post_konfig_resolve__(self, cfg: konfig.ConfigDict):
    self.cfg = cfg
```

--------------------------------

### Loading and running a serialized model

Source: https://github.com/google-research/kauldron/blob/main/docs/export.md

How to deserialize a model exported by Kauldron and run inference using JAX and Orbax. Note that the model must be run with the same number of devices as it was exported.

```python
import jax.export
from etils import epath
import orbax.checkpoint as ocp

WORKDIR = epath.Path("...")

PATH = WORKDIR / "train_model.jax_exported"
# Load the model
model = jax.export.deserialize(PATH.read_bytes())
# Load the params from the checkpoint
checkpointer = ocp.CheckpointManager(
  WORKDIR / "checkpoints",
  options=ocp.CheckpointManagerOptions(step_prefix="ckpt")
)
step = checkpointer.latest_step()
state = checkpointer.restore(step=step, args=ocp.args.StandardRestore())
params = state['params']

# Define a custom forward function that returns only the preds
# Here assuming the model inputs are called `image`
@jax.jit
def forward(image, params, key=jax.random.PRNGKey(0)):
  out = model.call(params=params, key=key, image=image)
  return out["preds"]


# Run the model on a custom input
forward(jnp.zeros((8, 64, 64, 3)), params)
```

--------------------------------

### Kauldron CLI: Data Element Spec (CLI)

Source: https://github.com/google-research/kauldron/blob/main/kauldron/cli/README.md

Use this command to display the element spec of the training data pipeline via the command line. It requires the path to the configuration file.

```sh
kauldron data element_spec \
    --cfg //third_party/py/kauldron/examples/mnist_autoencoder.py
```

--------------------------------

### Resolve ConfigDict to Python object

Source: https://github.com/google-research/kauldron/blob/main/kauldron/konfig/docs/demo.ipynb

Shows how to convert a `ConfigDict` back into its corresponding Python object using `konfig.resolve()`.

```python
model = konfig.resolve(model_cfg);
```

--------------------------------

### Define model inputs with kontext.Key

Source: https://github.com/google-research/kauldron/blob/main/kauldron/kontext/README.md

Use `kontext.Key` to define expected inputs for a model, allowing for flexible mapping to batch data.

```python
class MyModel(nn.Module):
  img: kontext.Key = kontext.REQUIRED  # Match `__call__` signature
  label: kontext.Key = kontext.REQUIRED

  @nn.compact
  def __call__(self, img, label):
    ...
```

--------------------------------

### Implement a custom data source

Source: https://github.com/google-research/kauldron/blob/main/kauldron/data/py/README.md

Inherit from DataSourceBase to wrap a grain.RandomAccessDataSource for cleaner configuration.

```python
class Tfds(kd.data.py.DataSourceBase):
  name: str
  split: str

  @functools.cached_property
  def data_source(self) -> grain.RandomAccessDataSource:
    return tfds.data_source(self.name, split=self.split)
```

--------------------------------

### Configure Metrics in Kauldron

Source: https://github.com/google-research/kauldron/blob/main/docs/metrics.md

Specify metrics and their inputs within the Kauldron configuration. This maps metric names to their implementations and data sources.

```python
cfg.metrics = {
    'reconstruction': kd.losses.L2(preds="preds.image", targets="batch.image"),
    'roc_auc': kd.metrics.RocAuc(preds="preds.logits", targets="batch.label"),
}
```

--------------------------------

### Wrap NNX Module with linen_from_nnx

Source: https://github.com/google-research/kauldron/blob/main/docs/nnx.md

Use `kd.contrib.nn.linen_from_nnx` to wrap an NNX module, making it compatible with Kauldron's Linen-based components. Constructor arguments can be passed directly.

```python
cfg.model = kd.contrib.nn.linen_from_nnx(MyModel)
```

```python
cfg.model = kd_nn.linen_from_nnx(nnx.Linear, in_features=3, out_features=4)
```

--------------------------------

### Map model inputs to batch data using string paths

Source: https://github.com/google-research/kauldron/blob/main/kauldron/kontext/README.md

Assign string paths to model `Key` attributes to specify how they should be extracted from a data batch.

```python
model = MyModel(
  img='batch.image',
  label='batch.label',
)
```

--------------------------------

### Define Data Pipeline Transforms

Source: https://github.com/google-research/kauldron/blob/main/docs/konfig_philosophy.md

Demonstrates low-level data pipeline configuration using a list of transforms.

```python
cfg = kd.data.tf.Tfds(
    name="ai2dcaption",
    split="train",
    shuffle=True,
    transforms=[
        # Low-level transforms
        gm.data.Tokenize(key="prompt", tokenizer=tokenizer, add_bos=True),
        gm.data.Tokenize(key="response", tokenizer=tokenizer, add_eos=True),
        gm.data.AddNextTokenPredictionFields(...),
        kd.data.Elements(keep=["input", "target", "loss_mask"]),
        gm.data.Pad(
            key=["input", "target", "loss_mask"],
            max_length=max_length,
            truncate=True,
        ),
        kd.data.Rearrange(key=["target", "loss_mask"], pattern="... -> ... 1"),
    ],
)
```

```python
cfg = kd.data.tf.Tfds(
    name="ai2dcaption",
    split="train",
    shuffle=True,
    transforms=[
        # High-level transform
        gm.data.Seq2SeqTask(in_prompt="prompt", in_response="response", ...),
    ],
)
```

--------------------------------

### Kauldron CLI: Data Element Spec (Python)

Source: https://github.com/google-research/kauldron/blob/main/kauldron/cli/README.md

This Python code snippet demonstrates how to execute the 'data element_spec' command programmatically. Ensure the 'kauldron' library is imported as 'kd'.

```python
from kauldron import kd
kd.cli.data.ElementSpec(cfg).execute()
```

--------------------------------

### Define Experiment Configuration

Source: https://github.com/google-research/kauldron/blob/main/docs/intro.md

Uses the konfig context manager to build ConfigDict objects that can be mutated and resolved into actual Python objects.

```python
with konfig.imports():
  import optax  # This looks like optax, but instead is a ConfigDict builder


cfg = optax.adam(learning_rate=0.003)  # This create a ConfigDict object !!!

assert cfg == konfig.ConfigDict({  # The config is a simple nested dict
    '__qualname__': 'optax:adam',
    'learning_rate': 0.003,
})

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

optimizer = konfig.resolve(cfg)  # Resolve the actual object (here optax.adam)
```

--------------------------------

### Dynamically define keys using __kontext_keys__ protocol

Source: https://github.com/google-research/kauldron/blob/main/kauldron/kontext/README.md

Implement `__kontext_keys__` in a class to dynamically provide key paths, useful for propagating keys from nested objects.

```python
@dataclasses.dataclass
class B:
  inner_keyed_obj: A

  def __kontext_keys__(self) -> dict[str, str | None]:
    return kontext.get_keypaths(self.inner_keyed_obj)


b = B(inner_keyed_obj=A(x='a[0].b.inner', y='a[1].e'))


assert kontext.get_keypaths(b) == {
    'x': 'a[0].b.inner',
    'y': 'a[1].e',
}
```

--------------------------------

### Create a config using mock_modules

Source: https://github.com/google-research/kauldron/blob/main/kauldron/konfig/docs/demo.ipynb

Uses `konfig.mock_modules()` to turn a nested Python call into a `konfig.ConfigDict` object without immediate execution.

```python
with konfig.mock_modules():
  model_cfg = nn.Sequential(layers=[
      nn.Dense(features=32),
      nn.Dropout(rate=0.5),
  ]);
```

--------------------------------

### Define Keys and Context

Source: https://github.com/google-research/kauldron/blob/main/docs/intro.md

Maps model inputs and losses to data sources using string identifiers.

```python
cfg.model = AutoEncoder(input="batch.image")

cfg.train_loss = kd.losses.L2(
    preds="preds.image",
    targets="batch.image",
)
```

```python
for batch in ds:

  preds = model(input=batch['image'])

  loss = train_loss(preds=preds['image'], targets=batch['image'])
```

--------------------------------

### Manage shape scopes with kt.typechecked

Source: https://github.com/google-research/kauldron/blob/main/kauldron/ktyping/README.md

Demonstrates how scopes are opened and closed automatically for functions, and how isinstance interacts with the active scope.

```python
@kt.typechecked  # < Opens a new scope when fn is called
def fn(a: Float["a"]) -> Int["a"]:  # arguments checked against this scope
  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.
  # So 'c' from Int["a c"] is not remembered (difference to jaxtyping).
  if isinstance(c, Int["a c"]):
    ...
    d = fn(a[:3])  # calling the function again opens (and closes) a new scope

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

--------------------------------

### Ignore Keys during Partial Loading

Source: https://github.com/google-research/kauldron/blob/main/docs/checkpoint.md

Configure PartialKauldronLoader to ignore missing or extra keys in the checkpoint using global patterns.

```python
cfg.init_transform = kd.ckpts.PartialKauldronLoader(
workdir=/path/to/old/workdir/,
ignore_model_keys=('**.missing_keys_from_checkpoint',),
ignore_restored_keys=('**.extra_keys_in_checkpoint',)
)
```

--------------------------------

### Generate typed paths with kontext.path_builder_from

Source: https://github.com/google-research/kauldron/blob/main/kauldron/kontext/README.md

Use `kontext.path_builder_from` to create typed path objects, enabling auto-completion and type checking for path construction.

```python
@flax.struct.dataclass
class Batch:
  image: jnp.array
  label: jnp.array

batch = kontext.path_builder_from('batch', Batch)

model = MyModel(
  img=batch.image,  # < Auto-complete and attribute checking, rather than `str`
  label=batch.label,
)
```

--------------------------------

### Annotate Keys in Modules

Source: https://github.com/google-research/kauldron/blob/main/docs/intro.md

Uses kontext.Key annotations to define expected inputs for Flax modules.

```python
from kauldron import kontext

class AutoEncoder(flax.linen.Module):
  input: kontext.Key  # Key names match the `__call__` signature

  @nn.compact
  def __call__(self, input: Float["*b h w c"]) -> Float["*b h w c"]:
    ...
```

--------------------------------

### Implementing a Custom TFDataPipeline Source

Source: https://github.com/google-research/kauldron/blob/main/kauldron/data/tf/README.md

Abstract method for implementing a custom data source within a TFDataPipeline. Ensure proper sharding and deterministic shuffling.

```python
class TFDataPipeline(kd.data.Pipeline):

  @abc.abstractmethod
  def ds_for_current_process(self) -> tf.data.Dataset:
    ...
```

--------------------------------

### Konfig Mutability and Serialization

Source: https://github.com/google-research/kauldron/blob/main/docs/konfig_philosophy.md

Demonstrating that Konfig objects are mutable and serializable.

```python
cfg.learning_rate = 1e-4
```

```python
cfg.to_json() == {
    '__qualname__': 'optax:adam',
    'learning_rate': 1e-4,
}
```

--------------------------------

### Demonstrate non-greedy matching in jaxtyping

Source: https://github.com/google-research/kauldron/blob/main/kauldron/ktyping/README.md

Shows how jaxtyping uses greedy matching which can lead to type checking failures in union types.

```python
@jt.jaxtyped(typechecker=typeguard.typechecked)
def foo(
  x: jt.Float[jt.Array, "a"] | jt.Float[jt.Array, "b"],
  y: jt.Float[jt.Array, "a b"]
): ...

foo(f32[3], f32[7 3])  # fails because x is greedily matched with Float["a"]
```

--------------------------------

### Customize Subdirectory Formats

Source: https://github.com/google-research/kauldron/blob/main/kauldron/xm/README.md

Override default directory naming conventions using the subdir_format attribute.

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

--------------------------------

### Mutate ConfigDict objects

Source: https://github.com/google-research/kauldron/blob/main/kauldron/konfig/docs/demo.ipynb

Demonstrates how to modify function calls, arguments, and constants within a configuration object.

```python
# Create a dummy function to demonstrate how to overwrite configs
def some_fn(*args, **kwargs):
  print(f'some_fn called with {args}, {kwargs}')


def other_fn(*args, **kwargs):
  print(f'other_fn called with {args}, {kwargs}')


# In Colab, import `__main__` to turn the Colab itself into configurable
import __main__


with konfig.mock_modules():
  cfg = __main__.some_fn(dtype=np.int32);
```

```python
cfg.x = 1
cfg.y = 2
cfg
```

```python
cfg[0] = 'a'  # Append additional arguments
cfg[1] = 'b'
cfg[-1] = 'b_overwritten'
cfg
```

```python
f'Previous qualname: {cfg.__qualname__}';

cfg.__qualname__ = '__main__:other_fn'  # Use `other_fn()` instead of `some_fn()`
cfg
```

```python
cfg.dtype.__const__ = 'numpy:float64'
cfg
```

```python
konfig.resolve(cfg)
```

--------------------------------

### Create a dataset mixture

Source: https://github.com/google-research/kauldron/blob/main/kauldron/data/py/README.md

Combines multiple datasets with nested transforms and global random cropping.

```python
cfg.train_ds = kd.data.py.Mix(
    datasets=[
        kd.data.py.Tfds(
            name='cifar100',
            split='train',
            transforms=[
                kd.data.py.Elements(keep=["image", "label"]),
            ],
        ),
        kd.data.py.Tfds(
            name='imagenet2012',
            split='train',
            transforms=[
                kd.data.py.Elements(keep=["image", "label"]),
                kd.data.py.Resize(key='image', height=32, width=32),
            ],
        ),
    ],
    seed=0,
    batch_size=256,
    transforms=[
        kd.data.py.RandomCrop(shape=(15, 15, None)),
    ],
)
```