def __call__(self, x: named_axes.NamedArray) -> named_axes.NamedArray:
    """Runs the caching attention computation and update the K/V cache state.

    When called, ``self.kv_cache_end_index`` should be filled with
    a scalar integer identifying the current size of the cache (before inserting
    this token), and ``self.kv_cache`` should be a `LocalState` that contains
    the current state.

    Args:
      x: The input to the computation, which will be mapped to queries, keys,
        and values by the sublayers.

    Returns:
      The final output of the ``attn_value_to_output`` sublayer.
    """

    # Retrieve effectful inputs.
    kvc_end_index = self.kv_cache_end_index.ask()
    key_cache, value_cache = self.kv_cache.get()

    # Compute queries, keys, and values as normal.
    query = self.input_to_query(x)
    key = self.input_to_key(x)
    value = self.input_to_value(x)

    # Update the KV caches.
    new_key_cache = named_axes.nmap(jax.lax.dynamic_update_slice)(
        key_cache.untag(self.sequence_axis),
        key.untag(self.sequence_axis),
        (kvc_end_index,),
    ).tag(self.sequence_axis)
    new_value_cache = named_axes.nmap(jax.lax.dynamic_update_slice)(
        value_cache.untag(self.sequence_axis),
        value.untag(self.sequence_axis),
        (kvc_end_index,),
    ).tag(self.sequence_axis)
    self.kv_cache.set((new_key_cache, new_value_cache))

    # Run the rest on the updated KV caches.
    attn = self.query_key_to_attn((query, new_key_cache))
    output = self.attn_value_to_output((attn, new_value_cache))
    return output
Runs the caching attention computation and update the K/V cache state.

    When called, ``self.kv_cache_end_index`` should be filled with
    a scalar integer identifying the current size of the cache (before inserting
    this token), and ``self.kv_cache`` should be a `LocalState` that contains
    the current state.

    Args:
      x: The input to the computation, which will be mapped to queries, keys,
        and values by the sublayers.

    Returns:
      The final output of the ``attn_value_to_output`` sublayer.
    
penzai/deprecated/v1/nn/attention.py