def from_uncached(
      cls,
      uncached: model_parts.TransformerLM,
      cache_len: int,
      batch_axes: dict[str, int],
      pad_id: int = 0,
      variable_name_prefix: str = "sampler",
  ) -> KVCachingTransformerLM:
    """Transforms a `Transformer` into cached sampling mode.

    This constructor hot-swaps all `pz.nn.Attention` 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 `Transformer` 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.
      pad_id: ID for the padding token.
      variable_name_prefix: Prefix for cached sampling variable names.

    Returns:
      A KVCachingTransformer.
    """

    def _fix_attn_mask(masker):
      if masker.kv_positions_input_name != "token_positions":
        raise ValueError(
            "Could not automatically convert attention mask layer with"
            f" non-standard positions input name: {masker}"
        )
      return dataclasses.replace(
          masker, kv_positions_input_name="kv_token_positions"
      )

    cached_axes = {
        **batch_axes,
        **uncached.metadata.common_head_axes,
        "projection": uncached.metadata.projection_dim,
    }
    attn_sel = pz.select(uncached.body).at_instances_of(pz.nn.Attention)
    fixed_attns = {}
    for ix, (keypath, attn) in enumerate(attn_sel.selected_by_path.items()):
      attn_with_new_kv_positions = (
          pz.select(attn)
          .at_instances_of(
              pz.nn.ApplyCausalAttentionMask
              | pz.nn.ApplyCausalSlidingWindowAttentionMask
          )
          .apply(_fix_attn_mask)
      )
      fixed_attns[keypath] = pz.nn.KVCachingAttention.from_uncached(
          attn_with_new_kv_positions,
          cache_len=cache_len,
          cached_axes=cached_axes,
          cache_dtype=uncached.metadata.activation_dtype,
          sequence_axis="seq",
          cache_end_index_key="cache_end_index",
          cache_label=f"{variable_name_prefix}/cache_{ix}",
          layerstack_axes=pz.nn.layerstack_axes_from_keypath(keypath),
      )
    caching_body = attn_sel.set_by_path(fixed_attns)
    return cls(
        metadata=uncached.metadata,
        cache_len=cache_len,
        batch_axes=batch_axes,
        pad_id=pad_id,
        previous_tokens=pz.StateVariable(
            value=pz.nx.zeros(
                {**batch_axes, "seq": cache_len}, dtype=jnp.int32
            ),
            label=f"{variable_name_prefix}/previous_tokens",
        ),
        cache_end_index=pz.StateVariable(
            value=0,
            label=f"{variable_name_prefix}/cache_end_index",
        ),
        body=caching_body,
    )
Transforms a `Transformer` into cached sampling mode.

    This constructor hot-swaps all `pz.nn.Attention` 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 `Transformer` 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.
      pad_id: ID for the padding token.
      variable_name_prefix: Prefix for cached sampling variable names.

    Returns:
      A KVCachingTransformer.
    
penzai/models/transformer/sampling_mode.py