def from_uncached(
      cls,
      original: model_core.GemmaAttention,
      cache_len: int,
      cached_axes: dict[str, int],
      cache_dtype: jax.typing.DTypeLike = jnp.float32,
  ) -> GemmaKVCachingAttention:
    """Builds a caching attention from an uncached attention.

    Args:
      original: The original attention layer that this block should replace.
      cache_len: Length of the cache; used to populate the initial state.
      cached_axes: Axis names and sizes for all other axes of the key and value
        arrays (e.g. for batch, heads, and the projected embeddings). These are
        used to initialize the cache.
      cache_dtype: Dtype for the data to store in the cache. Should match the
        dtype of the key and value arrays.

    Returns:
      A `GemmaKVCachingAttention` instance that behaves like the original
      `Attention` layer, but updates key-value caches iteratively, using new
      side input and state effect requests.
    """
    return super().from_uncached(
        original=original,
        sequence_axis="seq",
        cache_len=cache_len,
        cached_axes=cached_axes,
        cache_end_index_tag="cache_end_index",
        state_category="kv_cache",
        cache_dtype=cache_dtype,
    )
Builds a caching attention from an uncached attention.

    Args:
      original: The original attention layer that this block should replace.
      cache_len: Length of the cache; used to populate the initial state.
      cached_axes: Axis names and sizes for all other axes of the key and value
        arrays (e.g. for batch, heads, and the projected embeddings). These are
        used to initialize the cache.
      cache_dtype: Dtype for the data to store in the cache. Should match the
        dtype of the key and value arrays.

    Returns:
      A `GemmaKVCachingAttention` instance that behaves like the original
      `Attention` layer, but updates key-value caches iteratively, using new
      side input and state effect requests.
    
penzai/deprecated/v1/example_models/gemma/sampling_mode.py