def from_basic_segments(cls, tokens: pz.nx.NamedArray) -> GemmaInputs:
    """Constructs a simple input structure for a batch of single segments.

    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, as expected by the Gemma model.

    Args:
      tokens: Sequence of tokens, as an integer named array with a "seq" axis
        and possibly batch axes, which starts with the beginning-of-sequence
        token. Each 1d vector along the "seq" axis should represent an unpadded
        sequence.

    Returns:
      A full input structure containing the provided tokens, along with a simple
      incrementing position array and a causal mask.
    """
    seq = tokens.named_shape["seq"]
    # Query tokens can attend to keys/values if the query position is larger.
    attention_mask = pz.nx.arange("seq", seq) >= pz.nx.arange("kv_seq", seq)
    return cls(
        tokens=tokens,
        positions=pz.nx.arange("seq", seq),
        attention_mask=attention_mask,
    )
Constructs a simple input structure for a batch of single segments.

    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, as expected by the Gemma model.

    Args:
      tokens: Sequence of tokens, as an integer named array with a "seq" axis
        and possibly batch axes, which starts with the beginning-of-sequence
        token. Each 1d vector along the "seq" axis should represent an unpadded
        sequence.

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