penzai.example_models.gemma.simple_decoding_loop.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[source]#

Runs temperature sampling in a Python for loop.

  • 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.


A named array of continuations of the prompt.