def __call__(
      self, inputs: GemmaKVCachingInputs
  ) -> tuple[pz.nx.NamedArray, GemmaKVCachingState]:
    """Processes a new subsequence of tokens and adds them to the K/V cache.

    Args:
      inputs: Structure of input arguments, containing tokens, segment
        positions, an attention mask, and the current sampling state.

    Returns:
      A tuple ``(outputs, new_sampling_state)``, whre ``outputs`` is the final
      matrix of logits from the embedding decoding layer, which (in the normal
      configuration) will have axes "seq" and "vocabulary", and
      ``new_sampling_state`` is the updated sampling state with the updated
      key-value caches.
    """
    outs, kv_caches = self.body((
        (
            (inputs.tokens, inputs.positions, inputs.attention_mask),
            inputs.sampling_state.cache_end_index,
        ),
        inputs.sampling_state.kv_caches,
    ))
    return outs, GemmaKVCachingState(
        cache_len=inputs.sampling_state.cache_len,
        batch_axes=inputs.sampling_state.batch_axes,
        kv_caches=kv_caches,
        cache_end_index=(
            inputs.sampling_state.cache_end_index
            + inputs.tokens.named_shape["seq"]
        ),
    )
Processes a new subsequence of tokens and adds them to the K/V cache.

    Args:
      inputs: Structure of input arguments, containing tokens, segment
        positions, an attention mask, and the current sampling state.

    Returns:
      A tuple ``(outputs, new_sampling_state)``, whre ``outputs`` is the final
      matrix of logits from the embedding decoding layer, which (in the normal
      configuration) will have axes "seq" and "vocabulary", and
      ``new_sampling_state`` is the updated sampling state with the updated
      key-value caches.
    
penzai/deprecated/v1/example_models/gemma/sampling_mode.py