def from_basic_subsegments(
      cls, tokens: pz.nx.NamedArray, sampling_state: GemmaKVCachingState
  ) -> GemmaKVCachingInputs:
    """Constructs a simple input structure for a batch of unpadded samples.

    This can be used to process inputs that do not need advanced position or
    attention mask handling, and which just consist of ordinary sequences that
    are not packed together or padded. It augments the tokens with a standard
    position array and causal attention mask, adjusted by the current cache
    offset.

    Args:
      tokens: Subsequence of tokens, as an integer named array with a "seq" axis
        and possibly batch axes. When pre-filling, the "seq" axis can be the
        length of the prompt. When sampling, the "seq" instance will usually
        have length 1.
      sampling_state: Current sampling state, containing key-value caches.

    Returns:
      A full input structure containing the provided tokens, along with a simple
      incrementing position array and a causal mask, offset by the current
      sampling state.
    """
    seq = tokens.named_shape["seq"]
    offset = sampling_state.cache_end_index
    positions = pz.nx.arange("seq", seq) + offset
    # Query tokens can attend to keys/values if the query position is larger,
    # taking into account the cache offset.
    attention_mask = positions >= pz.nx.arange(
        "kv_seq", sampling_state.cache_len
    )
    return cls(
        tokens=tokens,
        positions=positions,
        attention_mask=attention_mask,
        sampling_state=sampling_state,
    )
Constructs a simple input structure for a batch of unpadded samples.

    This can be used to process inputs that do not need advanced position or
    attention mask handling, and which just consist of ordinary sequences that
    are not packed together or padded. It augments the tokens with a standard
    position array and causal attention mask, adjusted by the current cache
    offset.

    Args:
      tokens: Subsequence of tokens, as an integer named array with a "seq" axis
        and possibly batch axes. When pre-filling, the "seq" axis can be the
        length of the prompt. When sampling, the "seq" instance will usually
        have length 1.
      sampling_state: Current sampling state, containing key-value caches.

    Returns:
      A full input structure containing the provided tokens, along with a simple
      incrementing position array and a causal mask, offset by the current
      sampling state.
    
penzai/deprecated/v1/example_models/gemma/sampling_mode.py