def prefill(
    model: sampling_mode.GemmaKVCachingTransformer,
    initial_cache_state: sampling_mode.GemmaKVCachingState,
    prompt: pz.nx.NamedArray,
    pad_id: int,
) -> tuple[pz.nx.NamedArray, SamplingState]:
  """Prefills the key-value caches based on a prompt.

  Args:
    model: The converted model we are running inference with.
    initial_cache_state: The initial cache state created while converting the
      model.
    prompt: A named array of prompt tokens. Must have a "seq" axis along which
      the tokens for each batch element are arranged. Should usually start with
      the beginning of sequence token. This function assumes (but does not
      check) that all non-padding tokens preceed all padding tokens.
    pad_id: Token ID that corresponds to padding.

  Returns:
    A tuple ``(next_log_probs, sampling_state)``, where ``next_log_probs`` are
    the log-probabilities of the next token to sample, and ``sampling_state`` is
    a state that can be passed to future sampling calls.
  """
  not_padding_mask = prompt != pad_id

  # Since padding tokens are at the end, the offset of the last non-padding
  # token is one before the number of non-padding tokens.
  last_non_padding_index = (
      not_padding_mask.astype(jnp.int32).untag("seq").sum() - 1
  )

  # Token positions are just offsets along the token axis, since we're assuming
  # that padding only appears at the end.
  query_positions = pz.nx.arange("seq", prompt.named_shape["seq"])
  # Tokens can attend to any kv-token position that they are after.
  key_value_positions = pz.nx.arange("kv_seq", initial_cache_state.cache_len)
  attention_mask = query_positions >= key_value_positions
  # Run the model:
  out_logits, new_cache_state = model(
      sampling_mode.GemmaKVCachingInputs(
          tokens=prompt,
          positions=query_positions,
          attention_mask=attention_mask,
          sampling_state=initial_cache_state,
      )
  )
  # Extract the log probs from the last non-padding index, since that determines
  # the probs for the next sampled token.
  final_token_log_probs = pz.nx.nmap(jax.nn.log_softmax)(
      out_logits[{"seq": last_non_padding_index}].untag("vocabulary")
  ).tag("vocabulary")
  # Pad out the prompt to form our initial state matrix.
  padded_prompt = pz.nx.nmap(jnp.pad)(
      prompt.untag("seq"),
      [(0, initial_cache_state.cache_len - prompt.named_shape["seq"])],
      mode="constant",
      constant_values=pad_id,
  ).tag("seq")

  return final_token_log_probs, SamplingState(
      kv_cache_state=new_cache_state,
      previous_tokens=padded_prompt,
      pad_id=pad_id,
  )
Prefills the key-value caches based on a prompt.

  Args:
    model: The converted model we are running inference with.
    initial_cache_state: The initial cache state created while converting the
      model.
    prompt: A named array of prompt tokens. Must have a "seq" axis along which
      the tokens for each batch element are arranged. Should usually start with
      the beginning of sequence token. This function assumes (but does not
      check) that all non-padding tokens preceed all padding tokens.
    pad_id: Token ID that corresponds to padding.

  Returns:
    A tuple ``(next_log_probs, sampling_state)``, where ``next_log_probs`` are
    the log-probabilities of the next token to sample, and ``sampling_state`` is
    a state that can be passed to future sampling calls.
  
penzai/deprecated/v1/example_models/gemma/simple_decoding_loop.py