def advance_one_token(
    model: sampling_mode.GemmaKVCachingTransformer,
    state: SamplingState,
    next_token: jax.Array,
) -> tuple[pz.nx.NamedArray, SamplingState]:
  """Advances a sampling state by one token.

  This can be used to feed new sampled tokens one-at-a-time through the model,
  producing new log-probs that can be used to sample new tokens.

  Args:
    model: The converted model we are running inference with.
    state: The current sampling state.
    next_token: The next token to feed. Should not have a "seq" axis.

  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.
  """
  # Update the written tokens.
  update_t = state.kv_cache_state.cache_end_index
  updated_tokens = pz.nx.nmap(lambda s, v: s.at[update_t].set(v))(
      state.previous_tokens.untag("seq"), next_token
  ).tag("seq")
  assert dict(updated_tokens.named_shape) == dict(
      state.previous_tokens.named_shape
  )
  # Query position and key position are both based on the number of preceding
  # padding tokens (after updating the written tokens).
  nonpad_mask = updated_tokens != state.pad_id
  nonpad_so_far_inclusive = pz.nx.nmap(jnp.cumsum)(
      nonpad_mask.untag("seq"), dtype=jnp.int32
  ).tag("seq")
  nonpad_so_far_exclusive = nonpad_so_far_inclusive - nonpad_mask.astype(
      jnp.int32
  )
  query_positions = nonpad_so_far_inclusive[{"seq": -1}]
  key_value_positions = nonpad_so_far_exclusive.untag("seq").tag("kv_seq")
  # Tokens can attend to any kv-token position that they are after, as long as
  # it was NOT padding.
  attention_mask = pz.nx.nmap(jnp.logical_and)(
      query_positions >= key_value_positions,
      nonpad_mask.untag("seq").tag("kv_seq"),
  )
  # Run and update just like before, but add a tokens axis:
  out_logits, new_cache_state = model(
      sampling_mode.GemmaKVCachingInputs(
          tokens=next_token[{"seq": None}],
          positions=query_positions[{"seq": None}],
          attention_mask=attention_mask[{"seq": None}],
          sampling_state=state.kv_cache_state,
      )
  )
  # Extract the log probs from this token.
  final_token_log_probs = pz.nx.nmap(jax.nn.log_softmax)(
      out_logits.untag("seq").squeeze(0).untag("vocabulary")
  ).tag("vocabulary")
  return final_token_log_probs, SamplingState(
      kv_cache_state=new_cache_state,
      previous_tokens=updated_tokens,
      pad_id=state.pad_id,
  )
Advances a sampling state by one token.

  This can be used to feed new sampled tokens one-at-a-time through the model,
  producing new log-probs that can be used to sample new tokens.

  Args:
    model: The converted model we are running inference with.
    state: The current sampling state.
    next_token: The next token to feed. Should not have a "seq" axis.

  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