def __call__(
      self, tokens: pz.nx.NamedArray, **extra_side_inputs: dict[Any, Any]
  ) -> pz.nx.NamedArray:
    """Processes a new subsequence of tokens and adds them to the K/V cache.

    When called, the internal variables tracking the key-value cache will be
    updated with the new state.

    Args:
      tokens: Array of token IDs, as an integer named array with a "seq" axis
        and possibly batch axes. The batch axes must match the `batch_axes`
        attribute. Padding tokens are ignored.
      **extra_side_inputs: Extra side inputs, which will be forwarded on to the
        body. The "token_positions", "kv_token_positions", and "cache_end_index"
        inputs will be added automatically and do not need to be provided.

    Returns:
      Matrix of logits from the embedding decoding layer, which (in the
      normal configuration) will have axes "seq" and "vocabulary".
    """
    pz.chk.check_structure(
        tokens,
        pz.chk.ArraySpec(
            named_shape={**self.batch_axes, "seq": pz.chk.var("S")}
        ),
    )

    # Update the written tokens.
    updated_tokens = pz.nx.nmap(jax.lax.dynamic_update_slice)(
        self.previous_tokens.value.untag("seq"),
        tokens.untag("seq"),
        (self.cache_end_index.value,),
    ).tag("seq")
    # Query position and key position are both based on the number of preceding
    # padding tokens (after updating the written tokens). Compute by summing
    # up non-padding tokens, then slicing the part we are updating.
    kv_nonpad_mask = updated_tokens != self.pad_id
    kv_nonpad_so_far_inclusive = pz.nx.nmap(jnp.cumsum)(
        kv_nonpad_mask.untag("seq"), dtype=jnp.int32
    ).tag("seq")
    key_value_positions = pz.nx.nmap(jnp.where)(
        kv_nonpad_mask, kv_nonpad_so_far_inclusive - 1, -1
    )
    query_positions = pz.nx.nmap(jax.lax.dynamic_slice)(
        key_value_positions.untag("seq"),
        (self.cache_end_index.value,),
        (tokens.named_shape["seq"],),
    ).tag("seq")
    # Run the model.
    outs = self.body(
        tokens,
        token_positions=query_positions,
        kv_token_positions=key_value_positions,
        cache_end_index=self.cache_end_index.value,
        **extra_side_inputs,
    )
    # Update the state variables.
    self.previous_tokens.value = updated_tokens
    self.cache_end_index.value = (
        self.cache_end_index.value + tokens.named_shape["seq"]
    )

    return outs
Processes a new subsequence of tokens and adds them to the K/V cache.

    When called, the internal variables tracking the key-value cache will be
    updated with the new state.

    Args:
      tokens: Array of token IDs, as an integer named array with a "seq" axis
        and possibly batch axes. The batch axes must match the `batch_axes`
        attribute. Padding tokens are ignored.
      **extra_side_inputs: Extra side inputs, which will be forwarded on to the
        body. The "token_positions", "kv_token_positions", and "cache_end_index"
        inputs will be added automatically and do not need to be provided.

    Returns:
      Matrix of logits from the embedding decoding layer, which (in the
      normal configuration) will have axes "seq" and "vocabulary".
    
penzai/models/transformer/sampling_mode.py