GemmaKVCachingState

GemmaKVCachingState#

class penzai.deprecated.v1.example_models.gemma.sampling_mode.GemmaKVCachingState[source]#

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