def from_uncached(
      cls,
      original: Attention,
      sequence_axis: str,
      cache_len: int,
      cached_axes: dict[str, int],
      cache_end_index_tag: side_input.Tag,
      state_category: local_state.Category,
      cache_dtype: jax.typing.DTypeLike = jnp.float32,
  ) -> KVCachingAttention:
    """Builds a caching attention from an uncached attention.

    Args:
      original: The original attention layer that this block should replace.
      sequence_axis: The axis along which keys and values should be cached.
        Should be present in the output of the ``input_to_key`` and
        ``input_to_value`` sublayers.
      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_end_index_tag: Side input tag for the cache position side input.
        This should be used to identify the side inputs that should receive the
        cache position information, and should (usually) be provided to the
        `pz.de.WithSideInputsFromInputTuple` handler that actually provides this
        side input.
      state_category: Category for the local state. This should be used to
        identify the state variables that correspond to key-value caches in the
        model, and should (usually) be provided to the
        `pz.de.handle_local_states` call that functionalizes the state effect.
      cache_dtype: Dtype for the data to store in the cache. Should match the
        dtype of the key and value arrays.

    Returns:
      A ``KVCachingAttention`` instance that behaves like the original
      `Attention` layer, but updates key-value caches iteratively, using new
      side input and state effect requests.
    """

    def kv_cache_initializer():
      empty_cache = named_axes.zeros(
          {**cached_axes, sequence_axis: cache_len},
          dtype=cache_dtype,
      )
      return (empty_cache, empty_cache)

    return cls(
        input_to_query=original.input_to_query,
        input_to_key=original.input_to_key,
        input_to_value=original.input_to_value,
        query_key_to_attn=original.query_key_to_attn,
        attn_value_to_output=original.attn_value_to_output,
        sequence_axis=sequence_axis,
        kv_cache_end_index=side_input.SideInputRequest(cache_end_index_tag),
        kv_cache=local_state.InitialLocalStateRequest(
            kv_cache_initializer, category=state_category
        ),
    )
Builds a caching attention from an uncached attention.

    Args:
      original: The original attention layer that this block should replace.
      sequence_axis: The axis along which keys and values should be cached.
        Should be present in the output of the ``input_to_key`` and
        ``input_to_value`` sublayers.
      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_end_index_tag: Side input tag for the cache position side input.
        This should be used to identify the side inputs that should receive the
        cache position information, and should (usually) be provided to the
        `pz.de.WithSideInputsFromInputTuple` handler that actually provides this
        side input.
      state_category: Category for the local state. This should be used to
        identify the state variables that correspond to key-value caches in the
        model, and should (usually) be provided to the
        `pz.de.handle_local_states` call that functionalizes the state effect.
      cache_dtype: Dtype for the data to store in the cache. Should match the
        dtype of the key and value arrays.

    Returns:
      A ``KVCachingAttention`` 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/nn/attention.py