def temperature_sample_pyloop(
    model: sampling_mode.GemmaKVCachingTransformer,
    initial_cache_state: sampling_mode.GemmaKVCachingState,
    prompt: pz.nx.NamedArray,
    rng: jax.Array,
    pad_id: int,
    temperature: float = 1.0,
    max_sampling_steps: int | None = None,
) -> pz.nx.NamedArray:
  """Runs temperature sampling in a Python for loop.

  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.
    rng: JAX PRNGKey to use for sampling.
    pad_id: Token ID that corresponds to padding.
    temperature: Temperature to sample at.
    max_sampling_steps: Maximum number of sampling steps to run. If None,
      samples until filling up the key-value cache.

  Returns:
    A named array of continuations of the prompt.
  """
  if max_sampling_steps is None:
    max_sampling_steps = initial_cache_state.cache_len
  next_log_probs, sampling_state = prefill(
      model, initial_cache_state, prompt, pad_id
  )
  sampled_count = 0
  initial_cache_index = sampling_state.kv_cache_state.cache_end_index
  while True:
    kvcstate = sampling_state.kv_cache_state
    rng, key = jax.random.split(rng)
    # Split a key across named axes:
    batched_keys = pz.nx.random_split(key, kvcstate.batch_axes)
    next_token = pz.nx.nmap(jax.random.categorical)(
        batched_keys, next_log_probs.untag("vocabulary") / temperature
    )
    sampled_count += 1
    # Are we done?
    if (
        kvcstate.cache_end_index >= kvcstate.cache_len
        or sampled_count >= max_sampling_steps
    ):
      break
    next_log_probs, sampling_state = advance_one_token(
        model, sampling_state, next_token
    )
  # Add the last token we sampled (which doesn't need to be run through the
  # model).
  start = initial_cache_index
  end = sampling_state.kv_cache_state.cache_end_index
  final_written = pz.nx.concatenate(
      [
          sampling_state.previous_tokens[{"seq": pz.slice[start:end]}],
          next_token[{"seq": None}],
      ],
      "seq",
  )
  return final_written  # pytype: disable=bad-return-type
Runs temperature sampling in a Python for loop.

  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.
    rng: JAX PRNGKey to use for sampling.
    pad_id: Token ID that corresponds to padding.
    temperature: Temperature to sample at.
    max_sampling_steps: Maximum number of sampling steps to run. If None,
      samples until filling up the key-value cache.

  Returns:
    A named array of continuations of the prompt.
  
penzai/deprecated/v1/example_models/gemma/simple_decoding_loop.py