GemmaKVCachingState

class penzai.example_models.gemma.sampling_mode.GemmaKVCachingState

Bases: Struct

Sampling state for the key-value-caching Gemma variant.

You should not usually need to construct this on your own. Instead, it will be returned by GemmaKVCachingTransformer.from_uncached and updated by GemmaKVCachingTransformer.__call__.

Variables:

• cache_len (int) – The length of the key-value caches along the “seq” axis.

• batch_axes (dict[str, int]) – Axis names and sizes for the batch axes in the key-value caches.

• kv_caches (dict[str, Any]) – A dictionary of key-value caches extracted from the model.

• cache_end_index (int | jax.Array) – The current end index of the KV caches, used as the offset at which new keys and values will be inserted.

Methods

__init__(cache_len, batch_axes, kv_caches, ...)

Attributes

cache_len
batch_axes
kv_caches
cache_end_index

Inherited Methods

(expand to view inherited methods)

attributes_dict()	Constructs a dictionary with all of the fields in the class.
from_attributes(**field_values)	Directly instantiates a struct given all of its fields.
key_for_field(field_name)	Generates a JAX PyTree key for a given field name.
select()	Wraps this struct in a selection, enabling functional-style mutations.
tree_flatten()	Flattens this tree node.
tree_flatten_with_keys()	Flattens this tree node with keys.
tree_unflatten(aux_data, children)	Unflattens this tree node.
treescope_color()	Computes a CSS color to display for this object in treescope.