def from_uncached(
      cls,
      uncached: model_core.GemmaTransformer,
      cache_len: int,
      batch_axes: dict[str, int],
  ) -> tuple[GemmaKVCachingTransformer, GemmaKVCachingState]:
    """Transforms a `GemmaTransformer` into cached sampling mode.

    This constructor hot-swaps all `model_core.GemmaAttention` layers in the
    original model to enable key-value caching, then installs new handlers to
    update their states appropriately. Note that any modifications to the
    uncached model will persist in the decoding mode.

    Args:
      uncached: The original `GemmaTransformer` model.
      cache_len: Maximum sequence length for the key/value caches.
      batch_axes: Names and sizes for the batch axes that will be used for
        sampling. Required for initializing the key/value caches.

    Returns:
      Tuple ``(sampler_model, initial_sampling_state)``, where ``sampler_model``
      is a `GemmaKVCachingTransformer`, and ``initial_sampling_state`` holds the
      initial empty key/value caches.
    """
    cached_axes = {
        **batch_axes,
        "projection": uncached.config.projection_dim,
    }
    if not uncached.config.single_kv_head:
      cached_axes["heads"] = uncached.config.num_heads
    caching_body = (
        pz.select(uncached.body)
        .at_instances_of(model_core.GemmaAttention)
        .apply(
            lambda attn: GemmaKVCachingAttention.from_uncached(
                attn,
                cache_len=cache_len,
                cached_axes=cached_axes,
                cache_dtype=uncached.config.activation_dtype,
            )
        )
    )
    handled_body, initial_state = pz.de.handle_local_states(
        pz.de.WithSideInputsFromInputTuple.handling(
            caching_body, tags=["cache_end_index"]
        ),
        category="kv_cache",
    )
    inference_model = cls(config=uncached.config, body=handled_body)
    sampling_state = GemmaKVCachingState(
        cache_len=cache_len,
        batch_axes=batch_axes,
        kv_caches=initial_state,
        cache_end_index=0,
    )
    return inference_model, sampling_state
Transforms a `GemmaTransformer` into cached sampling mode.

    This constructor hot-swaps all `model_core.GemmaAttention` layers in the
    original model to enable key-value caching, then installs new handlers to
    update their states appropriately. Note that any modifications to the
    uncached model will persist in the decoding mode.

    Args:
      uncached: The original `GemmaTransformer` model.
      cache_len: Maximum sequence length for the key/value caches.
      batch_axes: Names and sizes for the batch axes that will be used for
        sampling. Required for initializing the key/value caches.

    Returns:
      Tuple ``(sampler_model, initial_sampling_state)``, where ``sampler_model``
      is a `GemmaKVCachingTransformer`, and ``initial_sampling_state`` holds the
      initial empty key/value caches.
    
penzai/deprecated/v1/example_models/gemma/sampling_mode.py