def from_uncached(
      cls,
      original: Attention,
      sequence_axis: str,
      cache_len: int,
      cached_axes: dict[str, int],
      cache_end_index_key: Hashable,
      cache_dtype: jax.typing.DTypeLike = jnp.float32,
      cache_label: variables.VariableLabel | None = None,
      layerstack_axes: dict[named_axes.AxisName, int] | None = None,
  ) -> 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_key: Key to use for the cache position side input.
      cache_dtype: Dtype for the data to store in the cache. Should match the
        dtype of the key and value arrays.
      cache_label: Optional label for the KV cache variable.
      layerstack_axes: Stacked axes that are used inside a LayerStack
        combinator. Usually inferred from `pz.nn.layerstack_axes_from_keypath`.

    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.
    """
    if layerstack_axes:
      cache_metadata = {
          "layerstack_axes": {
              ax: layer_stack.LayerStackVarBehavior.PER_LAYER
              for ax in layerstack_axes.keys()
          }
      }
    else:
      cache_metadata = {}
      layerstack_axes = {}
    empty_cache = named_axes.zeros(
        {**cached_axes, **layerstack_axes, sequence_axis: cache_len},
        dtype=cache_dtype,
    )

    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_key=cache_end_index_key,
        kv_cache=variables.StateVariable(
            label=cache_label,
            value=(empty_cache, empty_cache),
            metadata=cache_metadata,
        ),
    )
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_key: Key to use for the cache position side input.
      cache_dtype: Dtype for the data to store in the cache. Should match the
        dtype of the key and value arrays.
      cache_label: Optional label for the KV cache variable.
      layerstack_axes: Stacked axes that are used inside a LayerStack
        combinator. Usually inferred from `pz.nn.layerstack_axes_from_keypath`.

    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/nn/attention.py