Index

_

__abs__() (penzai.core.named_axes.NamedArrayBase method) __add__() (penzai.core.named_axes.NamedArrayBase method) __and__() (penzai.core.named_axes.NamedArrayBase method) __array__() (penzai.core.named_axes.NamedArrayBase method) __call__() (penzai.core.layer.Layer method) (penzai.core.struct.AbstractStructMetaclass method) (penzai.example_models.gemma.model_core.GemmaTransformer method) (penzai.example_models.gemma.sampling_mode.GemmaKVCachingTransformer method) (penzai.nn.attention.ApplyAttentionMask method) (penzai.nn.attention.Attention method) (penzai.nn.attention.KVCachingAttention method) (penzai.nn.combinators.Residual method) (penzai.nn.dropout.StochasticDropout method) (penzai.nn.embeddings.EmbeddingDecode method) (penzai.nn.embeddings.EmbeddingLookup method) (penzai.nn.grouping.CheckedSequential method) (penzai.nn.grouping.CheckStructure method) (penzai.nn.grouping.Identity method) (penzai.nn.grouping.NamedGroup method) (penzai.nn.grouping.Sequential method) (penzai.nn.linear_and_affine.AddBias method) (penzai.nn.linear_and_affine.ConstantRescale method) (penzai.nn.linear_and_affine.Linear method) (penzai.nn.linear_and_affine.LinearOperatorWeightInitializer method) (penzai.nn.linear_and_affine.NamedEinsum method) (penzai.nn.standardization.RMSStandardize method) (penzai.nn.standardization.Standardize method) (penzai.toolshed.basic_training.LossFunction method) (penzai.toolshed.basic_training.TrainStepFunction method) (penzai.toolshed.unflaxify.InterceptedFlaxModuleMethod method) (penzai.treescope.arrayviz.array_autovisualizer.ArrayAutovisualizer method) (penzai.treescope.autovisualize.Autovisualizer method) __divmod__() (penzai.core.named_axes.NamedArrayBase method) __enter__() (penzai.core.random_stream.RandomStream method) __exit__() (penzai.core.random_stream.RandomStream method) __floordiv__() (penzai.core.named_axes.NamedArrayBase method) __ge__() (penzai.core.named_axes.NamedArrayBase method) __getattr__() (penzai.data_effects.effect_base.EffectRequest method) (penzai.data_effects.effect_base.HandledEffectRef method) (penzai.toolshed.auto_nmap.AutoNmapModuleWrapper method) __getitem__() (penzai.core.named_axes.NamedArrayBase method) __gt__() (penzai.core.named_axes.NamedArrayBase method)

__init__() (penzai.core.context.ContextualValue method) (penzai.core.layer.MisconfiguredLayerError method) (penzai.core.shapecheck.StructureMismatchError method) (penzai.core.struct.PyTreeDataclassSafetyError method) (penzai.data_effects.effect_base.UnhandledEffectError method) (penzai.nn.parameters.UninitializedParameter method) (penzai.nn.parameters.UninitializedParameterError method) __init_subclass__() (penzai.core.layer.Layer class method) __invert__() (penzai.core.named_axes.NamedArrayBase method) __jax_array__() (penzai.core.named_axes.NamedArrayBase method) __le__() (penzai.core.named_axes.NamedArrayBase method) __len__() (penzai.core.selectors.Selection method) __lshift__() (penzai.core.named_axes.NamedArrayBase method) __lt__() (penzai.core.named_axes.NamedArrayBase method) __mod__() (penzai.core.named_axes.NamedArrayBase method) __mul__() (penzai.core.named_axes.NamedArrayBase method) __ne__() (penzai.core.named_axes.NamedArrayBase method) __neg__() (penzai.core.named_axes.NamedArrayBase method) __new__() (penzai.core.layer.MisconfiguredLayerError method) (penzai.core.shapecheck.StructureMismatchError method) (penzai.core.struct.PyTreeDataclassSafetyError method) (penzai.data_effects.effect_base.UnhandledEffectError method) (penzai.nn.parameters.UninitializedParameterError method) __or__() (penzai.core.named_axes.NamedArrayBase method) __pos__() (penzai.core.named_axes.NamedArrayBase method) __pow__() (penzai.core.named_axes.NamedArrayBase method) __radd__() (penzai.core.named_axes.NamedArrayBase method) __rand__() (penzai.core.named_axes.NamedArrayBase method) __rdivmod__() (penzai.core.named_axes.NamedArrayBase method) __rfloordiv__() (penzai.core.named_axes.NamedArrayBase method) __rlshift__() (penzai.core.named_axes.NamedArrayBase method) __rmod__() (penzai.core.named_axes.NamedArrayBase method) __rmul__() (penzai.core.named_axes.NamedArrayBase method) __ror__() (penzai.core.named_axes.NamedArrayBase method) __rpow__() (penzai.core.named_axes.NamedArrayBase method) __rrshift__() (penzai.core.named_axes.NamedArrayBase method) __rshift__() (penzai.core.named_axes.NamedArrayBase method) __rsub__() (penzai.core.named_axes.NamedArrayBase method) __rtruediv__() (penzai.core.named_axes.NamedArrayBase method) __rxor__() (penzai.core.named_axes.NamedArrayBase method) __sub__() (penzai.core.named_axes.NamedArrayBase method) __truediv__() (penzai.core.named_axes.NamedArrayBase method) __xor__() (penzai.core.named_axes.NamedArrayBase method)

A

AbstractStructMetaclass (class in penzai.core.struct) active_autovisualizer (in module penzai.treescope.autovisualize) active_expansion_strategy (in module penzai.treescope.default_renderer) active_renderer (in module penzai.treescope.default_renderer) add_parameter_prefix() (in module penzai.nn.parameters) AddBias (class in penzai.nn.linear_and_affine) advance_one_token() (in module penzai.example_models.gemma.simple_decoding_loop) Affine (class in penzai.nn.linear_and_affine) all() (penzai.core.named_axes.NamedArrayBase method) all_handler_ids() (in module penzai.data_effects.effect_base) AllowWordWrap (class in penzai.treescope.figures) annotate_shapes() (in module penzai.toolshed.annotate_shapes) any() (penzai.core.named_axes.NamedArrayBase method) apply() (penzai.core.selectors.Selection method) apply_and_inline() (penzai.core.selectors.Selection method) apply_with_selected_index() (penzai.core.selectors.Selection method) ApplyAttentionMask (class in penzai.nn.attention) ApplyRoPE (class in penzai.nn.embeddings) arange() (in module penzai.core.named_axes) argmax() (penzai.core.named_axes.NamedArrayBase method) argmin() (penzai.core.named_axes.NamedArrayBase method) argpartition() (penzai.core.named_axes.NamedArrayBase method) ArgsAndKwargs (class in penzai.toolshed.unflaxify) argsort() (penzai.core.named_axes.NamedArrayBase method) ArrayAutovisualizer (class in penzai.treescope.arrayviz.array_autovisualizer)

ArraySpec (class in penzai.core.shapecheck) ArrayvizDigitboxRendering (class in penzai.treescope.arrayviz.arrayviz) ArrayvizRendering (class in penzai.treescope.arrayviz.arrayviz) as_array_structure() (in module penzai.core.shapecheck) as_empty_parameter() (penzai.nn.parameters.UninitializedParameter method) as_frozen() (penzai.data_effects.local_state.InitialLocalStateRequest method) as_namedarrayview() (penzai.core.named_axes.NamedArrayBase method) ask() (penzai.data_effects.side_input.SideInputEffect method) assert_count_is() (penzai.core.selectors.Selection method) astype() (penzai.core.named_axes.NamedArrayBase method) at() (penzai.core.selectors.Selection method) at_childless() (penzai.core.selectors.Selection method) at_children() (penzai.core.selectors.Selection method) at_equal_to() (penzai.core.selectors.Selection method) at_instances_of() (penzai.core.selectors.Selection method) at_keypaths() (penzai.core.selectors.Selection method) at_pytree_leaves() (penzai.core.selectors.Selection method) at_subtrees_where() (penzai.core.selectors.Selection method) attach_shared_parameters() (in module penzai.nn.parameters) Attention (class in penzai.nn.attention) attributes_dict() (penzai.core.struct.Struct method) AutoNmapModuleWrapper (class in penzai.toolshed.auto_nmap) autovisualize() (penzai.treescope.treescope_ipython.AutovisualizerMagic method) Autovisualizer (class in penzai.treescope.autovisualize) AutovisualizerMagic (class in penzai.treescope.treescope_ipython)

B

basic_interactive_setup() (in module penzai.treescope.treescope_ipython) bolded() (in module penzai.treescope.figures) BranchAndAddTogether (class in penzai.nn.combinators) BranchAndMultiplyTogether (class in penzai.nn.combinators)

broadcast_like() (penzai.core.named_axes.NamedArrayBase method) broadcast_to() (penzai.core.named_axes.NamedArrayBase method) broken_handler_refs() (in module penzai.data_effects.effect_base) build_foldable_representation() (in module penzai.treescope.default_renderer) build_train_step_fn() (in module penzai.toolshed.basic_training)

C

call_and_extract_submodel() (in module penzai.toolshed.isolate_submodel) CalledWithManyStructures (class in penzai.toolshed.annotate_shapes) canonicalize() (penzai.core.named_axes.NamedArrayBase method) capture() (penzai.toolshed.unflaxify.ArgsAndKwargs class method) CastToDType (class in penzai.nn.basic_ops) check_layer() (in module penzai.toolshed.check_layers_by_tracing) check_no_duplicated_parameters() (in module penzai.nn.parameters) check_structure() (in module penzai.core.shapecheck) check_unique_axis_names_for_initializer() (in module penzai.nn.linear_and_affine) check_valid() (penzai.core.named_axes.NamedArrayBase method) checked_layer_call() (in module penzai.core.layer) CheckedSequential (class in penzai.nn.grouping) CheckStructure (class in penzai.nn.grouping) ChildAutovisualizer (class in penzai.treescope.autovisualize) choose() (penzai.core.named_axes.NamedArrayBase method) clip() (penzai.core.named_axes.NamedArrayBase method) CollectingSideOutputs (class in penzai.data_effects.side_output) color_from_string() (in module penzai.core.formatting_util)

combine() (in module penzai.core.partitioning) compress() (penzai.core.named_axes.NamedArrayBase method) compute_training_outputs_and_updates() (in module penzai.toolshed.basic_training) concatenate() (in module penzai.core.named_axes) conj() (penzai.core.named_axes.NamedArrayBase method) conjugate() (penzai.core.named_axes.NamedArrayBase method) constant_initializer() (in module penzai.nn.linear_and_affine) ConstantRescale (class in penzai.nn.linear_and_affine) ConstrainSharding (class in penzai.toolshed.sharding_util) ConstrainShardingByName (class in penzai.toolshed.sharding_util) ContextManagerMagic (class in penzai.treescope.treescope_ipython) ContextualValue (class in penzai.core.context) contract() (in module penzai.nn.linear_and_affine) count() (penzai.core.selectors.Selection method) CSSStyled (class in penzai.treescope.figures) cumprod() (penzai.core.named_axes.NamedArrayBase method) cumsum() (penzai.core.named_axes.NamedArrayBase method) CustomTreescopeVisualization (class in penzai.treescope.autovisualize)

D

dataclass_from_attributes() (in module penzai.core.dataclass_util) default_diverging_colormap (in module penzai.treescope.arrayviz.arrayviz) default_magic_autovisualizer (in module penzai.treescope.treescope_ipython) default_sequential_colormap (in module penzai.treescope.arrayviz.arrayviz) deselect() (penzai.core.selectors.Selection method) diagonal() (penzai.core.named_axes.NamedArrayBase method) DimensionVariableSubstitution (class in penzai.core.shapecheck) DimVar (class in penzai.core.shapecheck)

disable() (penzai.nn.dropout.StochasticDropout method) disable_interactive_context() (in module penzai.core.context) DisabledDropout (class in penzai.nn.dropout) display() (in module penzai.treescope.treescope_ipython) displayhook_with_ipython() (in module penzai.toolshed.patch_ipdb) dot() (penzai.core.named_axes.NamedArrayBase method) DropoutMLP (class in penzai.example_models.simple_mlp) dtype (penzai.core.named_axes.NamedArrayBase property) (penzai.core.shapecheck.ArraySpec attribute)

E

effect_protocol() (penzai.data_effects.effect_base.EffectHandler class method) (penzai.data_effects.effect_base.EffectRequest class method) (penzai.data_effects.effect_base.EffectRuntimeImpl class method) (penzai.data_effects.effect_base.HandledEffectRef class method) EffectHandler (class in penzai.data_effects.effect_base) EffectRequest (class in penzai.data_effects.effect_base) EffectRuntimeImpl (class in penzai.data_effects.effect_base)

Elementwise (class in penzai.nn.basic_ops) embed_shared_state_requests() (in module penzai.data_effects.local_state) EmbeddingDecode (class in penzai.nn.embeddings) EmbeddingLookup (class in penzai.nn.embeddings) EmbeddingTable (class in penzai.nn.embeddings) enable() (penzai.nn.dropout.DisabledDropout method) enable_interactive_context() (in module penzai.core.context)

F

flatten() (penzai.core.named_axes.NamedArrayBase method) flatten_selected_selections() (penzai.core.selectors.Selection method) floating_named() (penzai.core.shapecheck.ArraySpec class method) for_tokenizer() (penzai.treescope.arrayviz.array_autovisualizer.ArrayAutovisualizer class method) free_effect_types() (in module penzai.data_effects.effect_base) freeze_local_states() (in module penzai.data_effects.local_state) From (class in penzai.toolshed.model_rewiring) from_attributes() (penzai.core.struct.Struct class method) from_basic_segments() (penzai.example_models.gemma.model_core.GemmaInputs class method) from_basic_subsegments() (penzai.example_models.gemma.sampling_mode.GemmaKVCachingInputs class method) from_config() (penzai.data_effects.side_output.TellIntermediate class method) (penzai.example_models.gemma.model_core.GemmaAttention class method) (penzai.example_models.gemma.model_core.GemmaFeedForward class method) (penzai.example_models.gemma.model_core.GemmaTransformer class method) (penzai.example_models.gemma.model_core.GemmaTransformerBlock class method) (penzai.nn.attention.ApplyAttentionMask class method) (penzai.nn.embeddings.ApplyRoPE class method) (penzai.nn.embeddings.EmbeddingTable class method) (penzai.nn.linear_and_affine.AddBias class method) (penzai.nn.linear_and_affine.Linear class method) (penzai.nn.standardization.LayerNorm class method) (penzai.nn.standardization.RMSLayerNorm class method)

from_linear() (penzai.toolshed.lora.LowRankAdapter class method) from_object() (penzai.treescope.copypaste_fallback.NotRoundtrippable class method) from_pretrained() (penzai.example_models.gemma.model_core.GemmaTransformer class method) from_uncached() (penzai.example_models.gemma.sampling_mode.GemmaKVCachingAttention class method) (penzai.example_models.gemma.sampling_mode.GemmaKVCachingTransformer class method) (penzai.nn.attention.KVCachingAttention class method) from_uninitialized() (penzai.nn.parameters.ShareableUninitializedParameter class method) FrozenLocalStateRequest (class in penzai.data_effects.local_state) FrozenParameter (class in penzai.nn.parameters) full() (in module penzai.core.named_axes) full_substitute_dimension_variables() (in module penzai.core.shapecheck)

G

GemmaAttention (class in penzai.example_models.gemma.model_core) GemmaFeedForward (class in penzai.example_models.gemma.model_core) GemmaInputs (class in penzai.example_models.gemma.model_core) GemmaKVCachingAttention (class in penzai.example_models.gemma.sampling_mode) GemmaKVCachingInputs (class in penzai.example_models.gemma.sampling_mode) GemmaKVCachingState (class in penzai.example_models.gemma.sampling_mode) GemmaKVCachingTransformer (class in penzai.example_models.gemma.sampling_mode) GemmaTransformer (class in penzai.example_models.gemma.model_core) GemmaTransformerBlock (class in penzai.example_models.gemma.model_core)

GemmaTransformerConfig (class in penzai.example_models.gemma.model_core) get() (penzai.core.context.ContextualValue method) (penzai.core.selectors.Selection method) (penzai.data_effects.local_state.LocalStateEffect method) get_by_path() (penzai.core.selectors.Selection method) get_dimension_variables() (in module penzai.core.shapecheck) get_effect_color() (in module penzai.data_effects.effect_base) get_keypaths() (penzai.core.selectors.Selection method) get_sequence() (penzai.core.selectors.Selection method) GetPieceSize() (penzai.toolshed.token_visualization.SentencePieceProcessorLike method)

H

handle_local_states() (in module penzai.data_effects.local_state) HandledEffectRef (class in penzai.data_effects.effect_base) HandledLocalStateRef (class in penzai.data_effects.local_state) HandledRandomRef (class in penzai.data_effects.random) HandledSideInputRef (class in penzai.data_effects.side_input) HandledSideOutputRef (class in penzai.data_effects.side_output) handler_id() (penzai.data_effects.effect_base.EffectRuntimeImpl method) handling() (penzai.data_effects.random.WithFrozenRandomState class method) (penzai.data_effects.random.WithRandomKeyFromArg class method) (penzai.data_effects.random.WithStatefulRandomKey class method) (penzai.data_effects.side_input.WithConstantSideInputs class method) (penzai.data_effects.side_input.WithSideInputsFromInputTuple class method) (penzai.data_effects.side_output.CollectingSideOutputs class method)

hoist_constant_side_inputs() (in module penzai.data_effects.side_input) hoist_shared_state_requests() (in module penzai.data_effects.local_state) HoistedTag (class in penzai.data_effects.side_input)

I

Identity (class in penzai.nn.grouping) IdentityWithSavedActivations (class in penzai.toolshed.interleave_intermediates) IdToPiece() (penzai.toolshed.token_visualization.SentencePieceProcessorLike method) imag (penzai.core.named_axes.NamedArrayBase attribute) indented() (in module penzai.treescope.figures) infer_or_check_handler_id() (in module penzai.data_effects.effect_base) infer_rows_and_columns() (in module penzai.treescope.arrayviz.arrayviz) init_takes_fields() (in module penzai.core.dataclass_util) initial_state() (penzai.toolshed.basic_training.TrainState class method) initialize() (penzai.nn.parameters.UninitializedParameter method) initialize_parameters() (in module penzai.nn.parameters) initialize_parameters_sharded() (in module penzai.toolshed.sharding_util) initialize_with_value() (penzai.nn.parameters.UninitializedParameter method) InitialLocalStateRequest (class in penzai.data_effects.local_state) inline() (in module penzai.treescope.figures) inline_anonymous_sequentials() (in module penzai.nn.grouping) inline_groups() (in module penzai.nn.grouping) InlineBlock (class in penzai.treescope.figures) input_axes (penzai.nn.linear_and_affine.Linear property)

input_structure() (penzai.core.layer.Layer method) insert_after() (penzai.core.selectors.Selection method) insert_before() (penzai.core.selectors.Selection method) integer_digitbox() (in module penzai.treescope.arrayviz.arrayviz) InterceptedFlaxModuleMethod (class in penzai.toolshed.unflaxify) InterceptedFlaxScopeData (class in penzai.toolshed.unflaxify) into_pytree() (penzai.core.shapecheck.ArraySpec method) invert() (penzai.core.selectors.Selection method) IPythonVisualization (class in penzai.treescope.autovisualize) is_anonymous_sequential() (in module penzai.nn.grouping) is_empty() (penzai.core.selectors.Selection method) (penzai.core.shapecheck.DimensionVariableSubstitution method) is_namedarray() (in module penzai.core.named_axes) is_pytree_dataclass_type() (in module penzai.core.struct) is_pytree_node_field() (in module penzai.core.struct) is_sequential_or_named() (in module penzai.nn.grouping) IsControl() (penzai.toolshed.token_visualization.SentencePieceProcessorLike method) IsolatedSubmodel (class in penzai.toolshed.isolate_submodel) IsolationCapturer (class in penzai.toolshed.isolate_submodel) item() (penzai.core.named_axes.NamedArrayBase method)

J

Jitted (class in penzai.toolshed.jit_wrapper)

K

key_for_field() (penzai.core.struct.Struct method) KnockOutAttentionHeads (class in penzai.toolshed.model_rewiring)

KnownDim (class in penzai.core.shapecheck) KVCachingAttention (class in penzai.nn.attention)

L

Layer (class in penzai.core.layer) LayerNorm (class in penzai.nn.standardization) Linear (class in penzai.nn.linear_and_affine) LinearInPlace (class in penzai.nn.linear_and_affine) LinearizeAndAdjust (class in penzai.toolshed.model_rewiring) LinearOperatorWeightInitializer (class in penzai.nn.linear_and_affine)

load_arrayvis_javascript() (in module penzai.treescope.arrayviz.arrayviz) LocalStateEffect (class in penzai.data_effects.local_state) LocalStateEffectImpl (class in penzai.data_effects.local_state) loraify_linears_in_selection() (in module penzai.toolshed.lora) LossFunction (class in penzai.toolshed.basic_training) LowRankAdapter (class in penzai.toolshed.lora)

M

mark_shareable() (in module penzai.nn.parameters) MatchResult (class in penzai.core.shapecheck) max() (penzai.core.named_axes.NamedArrayBase method) maybe_dropout() (in module penzai.nn.dropout) mean() (penzai.core.named_axes.NamedArrayBase method) min() (penzai.core.named_axes.NamedArrayBase method) MisconfiguredLayerError MLP (class in penzai.example_models.simple_mlp) model (penzai.toolshed.basic_training.TrainState property) module penzai.core penzai.core.context penzai.core.dataclass_util penzai.core.formatting_util penzai.core.layer penzai.core.named_axes penzai.core.partitioning penzai.core.random_stream penzai.core.selectors penzai.core.shapecheck penzai.core.struct penzai.core.syntactic_sugar penzai.core.tree_util penzai.data_effects penzai.data_effects.effect_base penzai.data_effects.local_state penzai.data_effects.random penzai.data_effects.side_input penzai.data_effects.side_output penzai.example_models penzai.example_models.gemma penzai.example_models.gemma.model_core penzai.example_models.gemma.sampling_mode penzai.example_models.gemma.simple_decoding_loop penzai.example_models.simple_mlp penzai.nn penzai.nn.attention penzai.nn.basic_ops penzai.nn.combinators penzai.nn.dropout penzai.nn.embeddings penzai.nn.grouping penzai.nn.linear_and_affine penzai.nn.parameters penzai.nn.standardization penzai.pz penzai.pz.de penzai.pz.nn penzai.pz.ts penzai.toolshed penzai.toolshed.annotate_shapes penzai.toolshed.auto_nmap penzai.toolshed.basic_training penzai.toolshed.check_layers_by_tracing penzai.toolshed.interleave_intermediates penzai.toolshed.isolate_submodel penzai.toolshed.jit_wrapper penzai.toolshed.lora penzai.toolshed.model_rewiring penzai.toolshed.patch_ipdb penzai.toolshed.sharding_util penzai.toolshed.token_visualization penzai.toolshed.unflaxify penzai.treescope penzai.treescope.arrayviz penzai.treescope.arrayviz.array_autovisualizer penzai.treescope.arrayviz.arrayviz penzai.treescope.autovisualize penzai.treescope.copypaste_fallback penzai.treescope.default_renderer penzai.treescope.figures penzai.treescope.repr_lib penzai.treescope.treescope_ipython

mT (penzai.core.named_axes.NamedArrayBase attribute) MultiDimVar (class in penzai.core.shapecheck)

N

name_to_name_device_put() (in module penzai.toolshed.sharding_util) name_to_name_sharding() (in module penzai.toolshed.sharding_util) named_shape (penzai.core.named_axes.NamedArrayBase property) NamedArray (class in penzai.core.named_axes) NamedArrayBase (class in penzai.core.named_axes) NamedArrayView (class in penzai.core.named_axes) NamedEinsum (class in penzai.nn.linear_and_affine)

NamedGroup (class in penzai.nn.grouping) next_key() (penzai.core.random_stream.RandomStream method) (penzai.data_effects.random.RandomEffect method) nmap() (in module penzai.core.named_axes) nonzero() (penzai.core.named_axes.NamedArrayBase method) NotInThisPartition (class in penzai.core.partitioning) NotRoundtrippable (class in penzai.treescope.copypaste_fallback)

O

oklch_color() (in module penzai.core.formatting_util) ones() (in module penzai.core.named_axes) order_as() (penzai.core.named_axes.NamedArrayBase method)

order_like() (penzai.core.named_axes.NamedArrayBase method) output_axes (penzai.nn.linear_and_affine.Linear property) output_structure() (penzai.core.layer.Layer method)

P

parallel_axes (penzai.nn.linear_and_affine.Linear property) Parameter (class in penzai.nn.parameters) ParameterLike (class in penzai.nn.parameters) partition() (penzai.core.selectors.Selection method) patch_ipdb() (in module penzai.toolshed.patch_ipdb) path_matrix() (penzai.toolshed.model_rewiring.RewireComputationPaths method) penzai.core module penzai.core.context module penzai.core.dataclass_util module penzai.core.formatting_util module penzai.core.layer module penzai.core.named_axes module penzai.core.partitioning module penzai.core.random_stream module penzai.core.selectors module penzai.core.shapecheck module penzai.core.struct module penzai.core.syntactic_sugar module penzai.core.tree_util module penzai.data_effects module penzai.data_effects.effect_base module penzai.data_effects.local_state module penzai.data_effects.random module penzai.data_effects.side_input module penzai.data_effects.side_output module penzai.example_models module penzai.example_models.gemma module penzai.example_models.gemma.model_core module penzai.example_models.gemma.sampling_mode module penzai.example_models.gemma.simple_decoding_loop module penzai.example_models.simple_mlp module penzai.nn module penzai.nn.attention module penzai.nn.basic_ops module penzai.nn.combinators module penzai.nn.dropout module penzai.nn.embeddings module penzai.nn.grouping module

penzai.nn.linear_and_affine module penzai.nn.parameters module penzai.nn.standardization module penzai.pz module penzai.pz.de module penzai.pz.nn module penzai.pz.ts module penzai.toolshed module penzai.toolshed.annotate_shapes module penzai.toolshed.auto_nmap module penzai.toolshed.basic_training module penzai.toolshed.check_layers_by_tracing module penzai.toolshed.interleave_intermediates module penzai.toolshed.isolate_submodel module penzai.toolshed.jit_wrapper module penzai.toolshed.lora module penzai.toolshed.model_rewiring module penzai.toolshed.patch_ipdb module penzai.toolshed.sharding_util module penzai.toolshed.token_visualization module penzai.toolshed.unflaxify module penzai.treescope module penzai.treescope.arrayviz module penzai.treescope.arrayviz.array_autovisualizer module penzai.treescope.arrayviz.arrayviz module penzai.treescope.autovisualize module penzai.treescope.copypaste_fallback module penzai.treescope.default_renderer module penzai.treescope.figures module penzai.treescope.repr_lib module penzai.treescope.treescope_ipython module pick_nth_selected() (penzai.core.selectors.Selection method) positional_shape (penzai.core.named_axes.NamedArrayBase property) prefill() (in module penzai.example_models.gemma.simple_decoding_loop) pretty_keystr() (in module penzai.core.tree_util) PreventWordWrap (class in penzai.treescope.figures) prod() (penzai.core.named_axes.NamedArrayBase method) ptp() (penzai.core.named_axes.NamedArrayBase method) pytree_dataclass() (in module penzai.core.struct) PyTreeDataclassSafetyError

R

random_split() (in module penzai.core.named_axes) RandomEffect (class in penzai.data_effects.random) RandomEffectImpl (class in penzai.data_effects.random) RandomRequest (class in penzai.data_effects.random) RandomStream (class in penzai.core.random_stream) ravel() (penzai.core.named_axes.NamedArrayBase method) real (penzai.core.named_axes.NamedArrayBase attribute) refine() (penzai.core.selectors.Selection method) register_as_default() (in module penzai.treescope.treescope_ipython) register_autovisualize_magic() (in module penzai.treescope.treescope_ipython) register_context_manager_magic() (in module penzai.treescope.treescope_ipython) register_effect_color() (in module penzai.data_effects.effect_base) RemainingAxisPlaceholder (class in penzai.core.shapecheck) remove_from_parent() (penzai.core.selectors.Selection method) RenameAxes (class in penzai.nn.linear_and_affine) render_array() (in module penzai.treescope.arrayviz.arrayviz)

render_array_sharding() (in module penzai.treescope.arrayviz.arrayviz) render_dictionary_wrapper() (in module penzai.treescope.repr_lib) render_not_roundtrippable() (in module penzai.treescope.copypaste_fallback) render_object_constructor() (in module penzai.treescope.repr_lib) render_sharded_shape() (in module penzai.treescope.arrayviz.arrayviz) render_to_html() (in module penzai.treescope.default_renderer) render_to_text() (in module penzai.treescope.default_renderer) RendersAsRootInIPython (class in penzai.treescope.figures) repeat() (penzai.core.named_axes.NamedArrayBase method) reshape() (penzai.core.named_axes.NamedArrayBase method) Residual (class in penzai.nn.combinators) RewireComputationPaths (class in penzai.toolshed.model_rewiring) RMSLayerNorm (class in penzai.nn.standardization) RMSStandardize (class in penzai.nn.standardization) round() (penzai.core.named_axes.NamedArrayBase method) run_and_interleave_intermediates() (in module penzai.toolshed.interleave_intermediates)

S

SamplingState (class in penzai.example_models.gemma.simple_decoding_loop) searchsorted() (penzai.core.named_axes.NamedArrayBase method) select() (in module penzai.core.selectors) (penzai.core.struct.Struct method) select_and_set_by_path() (penzai.core.selectors.Selection method) Selection (class in penzai.core.selectors) SelectionHole (class in penzai.core.selectors) SelectionQuote (class in penzai.core.selectors) SentencePieceProcessorLike (class in penzai.toolshed.token_visualization) Sequential (class in penzai.nn.grouping) set() (penzai.core.selectors.Selection method) (penzai.data_effects.local_state.LocalStateEffect method) set_by_path() (penzai.core.selectors.Selection method) set_interactive() (penzai.core.context.ContextualValue method) set_scoped() (penzai.core.context.ContextualValue method) set_sequence() (penzai.core.selectors.Selection method) ShareableUninitializedParameter (class in penzai.nn.parameters) SharedLocalStateRequest (class in penzai.data_effects.local_state) SharedParameterLookup (class in penzai.nn.parameters) SharedParamTag (class in penzai.nn.parameters) show() (in module penzai.treescope.treescope_ipython) show_selection() (penzai.core.selectors.Selection method) show_token_array() (in module penzai.toolshed.token_visualization) show_token_scores() (in module penzai.toolshed.token_visualization)

show_value() (penzai.core.selectors.Selection method) SideInputEffect (class in penzai.data_effects.side_input) SideInputEffectImpl (class in penzai.data_effects.side_input) SideInputRequest (class in penzai.data_effects.side_input) SideOutputEffect (class in penzai.data_effects.side_output) SideOutputEffectImpl (class in penzai.data_effects.side_output) SideOutputRequest (class in penzai.data_effects.side_output) SideOutputValue (class in penzai.data_effects.side_output) SliceLike (class in penzai.core.syntactic_sugar) Softmax (class in penzai.nn.basic_ops) sort() (penzai.core.named_axes.NamedArrayBase method) squeeze() (penzai.core.named_axes.NamedArrayBase method) stack() (in module penzai.core.named_axes) Standardize (class in penzai.nn.standardization) Static (class in penzai.toolshed.annotate_shapes) std() (penzai.core.named_axes.NamedArrayBase method) StochasticDropout (class in penzai.nn.dropout) Struct (class in penzai.core.struct) StructStaticMetadata (class in penzai.core.struct) StructureMismatchError styled() (in module penzai.treescope.figures) sum() (penzai.core.named_axes.NamedArrayBase method) SupportsParameterRenaming (class in penzai.nn.parameters) swapaxes() (penzai.core.named_axes.NamedArrayBase method)

T

T (penzai.core.named_axes.NamedArrayBase attribute) tag() (penzai.core.named_axes.NamedArrayBase method) tag_prefix() (penzai.core.named_axes.NamedArrayBase method) TaggedRandomRequest (class in penzai.data_effects.random) take() (penzai.core.named_axes.NamedArrayBase method) tell() (penzai.data_effects.side_output.SideOutputEffect method) TellIntermediate (class in penzai.data_effects.side_output) temperature_sample_pyloop() (in module penzai.example_models.gemma.simple_decoding_loop) text_on_color() (in module penzai.treescope.arrayviz.arrayviz) TmpPosAxisMarker (class in penzai.core.named_axes)

trace() (penzai.core.named_axes.NamedArrayBase method) TrainState (class in penzai.toolshed.basic_training) TrainStepFunction (class in penzai.toolshed.basic_training) transpose() (penzai.core.named_axes.NamedArrayBase method) tree_flatten() (penzai.core.struct.Struct method) tree_flatten_exactly_one_level() (in module penzai.core.tree_util) tree_flatten_with_keys() (penzai.core.struct.Struct method) tree_unflatten() (penzai.core.struct.Struct class method) treescope_color() (penzai.core.struct.Struct method) TreescopeRenderingFigure (class in penzai.treescope.figures)

U

unchecked_layer_call() (in module penzai.core.layer) unflaxify_apply() (in module penzai.toolshed.unflaxify) UnhandledEffectError UninitializedParameter (class in penzai.nn.parameters) UninitializedParameterError unpatch_ipdb() (in module penzai.toolshed.patch_ipdb)

unsafe_mark_active() (penzai.core.random_stream.RandomStream method) unstack() (in module penzai.core.named_axes) untag() (penzai.core.named_axes.NamedArrayBase method) untag_prefix() (penzai.core.named_axes.NamedArrayBase method) unwrap() (penzai.core.named_axes.NamedArrayBase method) use_autovisualizer_if_present() (in module penzai.treescope.autovisualize) using_expansion_strategy() (in module penzai.treescope.default_renderer)

V

value (penzai.nn.parameters.ParameterLike property) (penzai.nn.parameters.SharedParameterLookup property) (penzai.nn.parameters.UninitializedParameter property) value_structure (penzai.nn.parameters.ParameterLike property) ValueColoredTextbox (class in penzai.treescope.arrayviz.arrayviz)

var() (in module penzai.core.shapecheck) (penzai.core.named_axes.NamedArrayBase method) variance_scaling_initializer() (in module penzai.nn.linear_and_affine) vars_for_axes() (in module penzai.core.shapecheck) view() (penzai.core.named_axes.NamedArrayBase method)

W

where() (penzai.core.selectors.Selection method) Wildcard (class in penzai.core.shapecheck) with_() (penzai.treescope.treescope_ipython.ContextManagerMagic method) with_color() (in module penzai.treescope.figures) with_font_size() (in module penzai.treescope.figures) with_positional_prefix() (penzai.core.named_axes.NamedArrayBase method) (penzai.core.named_axes.NamedArrayView method) with_renamed_parameters() (penzai.nn.parameters.SupportsParameterRenaming method) WithConstantSideInputs (class in penzai.data_effects.side_input)

WithFrozenRandomState (class in penzai.data_effects.random) WithFunctionalLocalState (class in penzai.data_effects.local_state) WithRandomKeyFromArg (class in penzai.data_effects.random) WithSideInputsFromInputTuple (class in penzai.data_effects.side_input) WithStatefulRandomKey (class in penzai.data_effects.random) wrap() (in module penzai.core.named_axes) (penzai.core.named_axes.NamedArray class method) wrap_as_treescope_figure() (in module penzai.treescope.figures) wrap_module() (in module penzai.toolshed.auto_nmap)

Z

zero_initializer() (in module penzai.nn.linear_and_affine)

zeros() (in module penzai.core.named_axes)