Index _ | A | B | C | D | E | F | G | H | I | J | K | L | M | N | O | P | R | S | T | U | V | W | Z _ __abs__() (penzai.core.named_axes.NamedArrayBase method) __add__() (penzai.core.named_axes.NamedArrayBase method) __and__() (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.experimental.v2.models.transformer.model_parts.TransformerLM method) (penzai.experimental.v2.models.transformer.sampling_mode.KVCachingTransformerLM method) (penzai.experimental.v2.nn.attention.ApplyCausalAttentionMask method) (penzai.experimental.v2.nn.attention.ApplyCausalSlidingWindowAttentionMask method) (penzai.experimental.v2.nn.attention.ApplyExplicitAttentionMask method) (penzai.experimental.v2.nn.attention.Attention method) (penzai.experimental.v2.nn.attention.KVCachingAttention method) (penzai.experimental.v2.nn.combinators.Residual method) (penzai.experimental.v2.nn.dropout.StochasticDropout method) (penzai.experimental.v2.nn.embeddings.EmbeddingDecode method) (penzai.experimental.v2.nn.embeddings.EmbeddingLookup method) (penzai.experimental.v2.nn.grouping.CheckedSequential method) (penzai.experimental.v2.nn.grouping.CheckStructure method) (penzai.experimental.v2.nn.grouping.Identity method) (penzai.experimental.v2.nn.grouping.NamedGroup method) (penzai.experimental.v2.nn.grouping.Sequential method) (penzai.experimental.v2.nn.layer.Layer method) (penzai.experimental.v2.nn.layer_stack.LayerStack method) (penzai.experimental.v2.nn.linear_and_affine.AddBias method) (penzai.experimental.v2.nn.linear_and_affine.ConstantRescale method) (penzai.experimental.v2.nn.linear_and_affine.Linear method) (penzai.experimental.v2.nn.linear_and_affine.LinearOperatorWeightInitializer method) (penzai.experimental.v2.nn.linear_and_affine.NamedEinsum method) (penzai.experimental.v2.nn.standardization.RMSStandardize method) (penzai.experimental.v2.nn.standardization.Standardize method) (penzai.experimental.v2.toolshed.basic_training.LossFunction method) (penzai.experimental.v2.toolshed.unflaxify.InterceptedFlaxModuleMethod 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.experimental.v2.core.variables.LabeledVariable method) (penzai.experimental.v2.core.variables.StateVariable method) (penzai.experimental.v2.core.variables.UnboundVariableError method) (penzai.experimental.v2.core.variables.VariableConflictError 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) __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.experimental.v2.core.variables.UnboundVariableError method) (penzai.experimental.v2.core.variables.VariableConflictError 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) AbstractVariable (class in penzai.experimental.v2.core.variables) AbstractVariableSlot (class in penzai.experimental.v2.core.variables) AbstractVariableValue (class in penzai.experimental.v2.core.variables) 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.experimental.v2.nn.linear_and_affine) (class in penzai.nn.linear_and_affine) advance_one_token() (in module penzai.example_models.gemma.simple_decoding_loop) Affine (class in penzai.experimental.v2.nn.linear_and_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) ApplyCausalAttentionMask (class in penzai.experimental.v2.nn.attention) ApplyCausalSlidingWindowAttentionMask (class in penzai.experimental.v2.nn.attention) ApplyExplicitAttentionMask (class in penzai.experimental.v2.nn.attention) ApplyRoPE (class in penzai.experimental.v2.nn.embeddings) (class in penzai.nn.embeddings) ApplyRoPEToSubset (class in penzai.experimental.v2.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.experimental.v2.toolshed.unflaxify) (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) assert_no_parameter_slots() (in module penzai.experimental.v2.nn.parameters) astype() (penzai.core.named_axes.NamedArrayBase method) at (penzai.core.named_axes.NamedArrayBase property) 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.experimental.v2.nn.attention) (class in penzai.nn.attention) AttentionTypeGlobalCausal (class in penzai.experimental.v2.models.transformer.variants.llamalike_common) AttentionTypeSlidingWindowCausal (class in penzai.experimental.v2.models.transformer.variants.llamalike_common) attributes_dict() (penzai.core.struct.Struct method) AutoNmapModuleWrapper (class in penzai.toolshed.auto_nmap) AutoOrderedAcrossTypes (class in penzai.experimental.v2.core.auto_order_types) AutoStateVarLabel (class in penzai.experimental.v2.core.variables) 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) bind_variables() (in module penzai.experimental.v2.core.variables) (penzai.experimental.v2.nn.layer.Layer method) bolded() (in module penzai.treescope.figures) BranchAndAddTogether (class in penzai.experimental.v2.nn.combinators) (class in penzai.nn.combinators) BranchAndMultiplyTogether (class in penzai.experimental.v2.nn.combinators) (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_gpt_neox_attention() (in module penzai.experimental.v2.models.transformer.variants.gpt_neox) build_gpt_neox_block() (in module penzai.experimental.v2.models.transformer.variants.gpt_neox) build_gpt_neox_feedforward() (in module penzai.experimental.v2.models.transformer.variants.gpt_neox) build_gpt_neox_transformer() (in module penzai.experimental.v2.models.transformer.variants.gpt_neox) build_llamalike_attention() (in module penzai.experimental.v2.models.transformer.variants.llamalike_common) build_llamalike_block() (in module penzai.experimental.v2.models.transformer.variants.llamalike_common) build_llamalike_feedforward() (in module penzai.experimental.v2.models.transformer.variants.llamalike_common) build_llamalike_transformer() (in module penzai.experimental.v2.models.transformer.variants.llamalike_common) build_train_step_fn() (in module penzai.toolshed.basic_training) C call_and_extract_submodel() (in module penzai.experimental.v2.toolshed.isolate_submodel) (in module penzai.toolshed.isolate_submodel) CalledWithManyStructures (class in penzai.toolshed.annotate_shapes) canonicalize() (penzai.core.named_axes.NamedArrayBase method) capture() (penzai.experimental.v2.toolshed.unflaxify.ArgsAndKwargs class method) (penzai.toolshed.unflaxify.ArgsAndKwargs class method) CastToDType (class in penzai.experimental.v2.nn.basic_ops) (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.experimental.v2.nn.linear_and_affine) (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.experimental.v2.nn.grouping) (class in penzai.nn.grouping) CheckStructure (class in penzai.experimental.v2.nn.grouping) (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.experimental.v2.nn.linear_and_affine) (in module penzai.nn.linear_and_affine) ConstantRescale (class in penzai.experimental.v2.nn.linear_and_affine) (class in penzai.nn.linear_and_affine) ConstrainSharding (class in penzai.experimental.v2.toolshed.sharding_util) (class in penzai.toolshed.sharding_util) ConstrainShardingByName (class in penzai.experimental.v2.toolshed.sharding_util) (class in penzai.toolshed.sharding_util) ContextManagerMagic (class in penzai.treescope.treescope_ipython) ContextualValue (class in penzai.core.context) contract() (in module penzai.experimental.v2.nn.linear_and_affine) (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) derive_param_key() (in module penzai.experimental.v2.nn.parameters) 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.experimental.v2.nn.dropout.StochasticDropout method) (penzai.nn.dropout.StochasticDropout method) disable_interactive_context() (in module penzai.core.context) DisabledDropout (class in penzai.experimental.v2.nn.dropout) (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) (class in penzai.experimental.v2.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.experimental.v2.nn.basic_ops) (class in penzai.nn.basic_ops) embed_shared_state_requests() (in module penzai.data_effects.local_state) EmbeddingDecode (class in penzai.experimental.v2.nn.embeddings) (class in penzai.nn.embeddings) EmbeddingLookup (class in penzai.experimental.v2.nn.embeddings) (class in penzai.nn.embeddings) EmbeddingTable (class in penzai.experimental.v2.nn.embeddings) (class in penzai.nn.embeddings) enable() (penzai.experimental.v2.nn.dropout.DisabledDropout method) (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) FlaxVarLabel (class in penzai.experimental.v2.toolshed.unflaxify) 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() (penzai.experimental.v2.core.variables.AbstractVariable method) freeze_local_states() (in module penzai.data_effects.local_state) freeze_params() (in module penzai.experimental.v2.core.variables) freeze_state_vars() (in module penzai.experimental.v2.core.variables) freeze_variables() (in module penzai.experimental.v2.core.variables) From (class in penzai.experimental.v2.toolshed.model_rewiring) (class in penzai.toolshed.model_rewiring) from_attributes() (penzai.core.struct.Struct class method) from_base_key() (penzai.experimental.v2.core.random_stream.RandomStream 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.experimental.v2.nn.embeddings.EmbeddingTable class method) (penzai.experimental.v2.nn.linear_and_affine.AddBias class method) (penzai.experimental.v2.nn.linear_and_affine.Linear class method) (penzai.experimental.v2.nn.standardization.LayerNorm class method) (penzai.experimental.v2.nn.standardization.RMSLayerNorm 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.experimental.v2.toolshed.lora.LowRankAdapter class method) (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_sublayer_builder() (penzai.experimental.v2.nn.layer_stack.LayerStack class method) from_uncached() (penzai.example_models.gemma.sampling_mode.GemmaKVCachingAttention class method) (penzai.example_models.gemma.sampling_mode.GemmaKVCachingTransformer class method) (penzai.experimental.v2.models.transformer.sampling_mode.KVCachingTransformerLM class method) (penzai.experimental.v2.nn.attention.KVCachingAttention 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 gemma_from_pretrained_checkpoint() (in module penzai.experimental.v2.models.transformer.variants.gemma) 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) get_slot() (penzai.experimental.v2.core.variables.AbstractVariable method) (penzai.experimental.v2.core.variables.AbstractVariableValue method) GetPieceSize() (penzai.toolshed.token_visualization.SentencePieceProcessorLike method) gpt_neox_from_huggingface_model() (in module penzai.experimental.v2.models.transformer.variants.gpt_neox) GPTNeoXTransformerConfig (class in penzai.experimental.v2.models.transformer.variants.gpt_neox) 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.experimental.v2.nn.grouping) (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.experimental.v2.nn.grouping) (in module penzai.nn.grouping) inline_groups() (in module penzai.experimental.v2.nn.grouping) (in module penzai.nn.grouping) InlineBlock (class in penzai.treescope.figures) input_axes (penzai.experimental.v2.nn.linear_and_affine.Linear property) (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.experimental.v2.toolshed.unflaxify) (class in penzai.toolshed.unflaxify) InterceptedFlaxScopeData (class in penzai.experimental.v2.toolshed.unflaxify) (class in penzai.toolshed.unflaxify) InternalTrainerState (class in penzai.experimental.v2.toolshed.basic_training) 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.experimental.v2.nn.grouping) (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.experimental.v2.nn.grouping) (in module penzai.nn.grouping) IsControl() (penzai.toolshed.token_visualization.SentencePieceProcessorLike method) IsolatedSubmodel (class in penzai.experimental.v2.toolshed.isolate_submodel) (class in penzai.toolshed.isolate_submodel) IsolationCapturer (class in penzai.experimental.v2.toolshed.isolate_submodel) (class in penzai.toolshed.isolate_submodel) item() (penzai.core.named_axes.NamedArrayBase method) J Jitted (class in penzai.experimental.v2.toolshed.jit_wrapper) (class in penzai.toolshed.jit_wrapper) K key_for_field() (penzai.core.struct.Struct method) (penzai.experimental.v2.nn.layer_stack.LayerStack method) KnockOutAttentionHeads (class in penzai.experimental.v2.toolshed.model_rewiring) (class in penzai.toolshed.model_rewiring) KnownDim (class in penzai.core.shapecheck) KVCachingAttention (class in penzai.experimental.v2.nn.attention) (class in penzai.nn.attention) KVCachingTransformerLM (class in penzai.experimental.v2.models.transformer.sampling_mode) L LabeledVariable (class in penzai.experimental.v2.core.variables) LabeledVariableValue (class in penzai.experimental.v2.core.variables) Layer (class in penzai.core.layer) (class in penzai.experimental.v2.nn.layer) LayerNorm (class in penzai.experimental.v2.nn.standardization) (class in penzai.nn.standardization) LayerStack (class in penzai.experimental.v2.nn.layer_stack) layerstack_axes_from_keypath() (in module penzai.experimental.v2.nn.layer_stack) LayerStackGetAttrKey (class in penzai.experimental.v2.nn.layer_stack) LayerStackVarBehavior (class in penzai.experimental.v2.nn.layer_stack) Linear (class in penzai.experimental.v2.nn.linear_and_affine) (class in penzai.nn.linear_and_affine) LinearInPlace (class in penzai.experimental.v2.nn.linear_and_affine) (class in penzai.nn.linear_and_affine) LinearizeAndAdjust (class in penzai.experimental.v2.toolshed.model_rewiring) (class in penzai.toolshed.model_rewiring) LinearOperatorWeightInitializer (class in penzai.experimental.v2.nn.linear_and_affine) (class in penzai.nn.linear_and_affine) llama_from_huggingface_model() (in module penzai.experimental.v2.models.transformer.variants.llama) llamalike_from_huggingface_model() (in module penzai.experimental.v2.models.transformer.variants.llamalike_common) LlamalikeTransformerConfig (class in penzai.experimental.v2.models.transformer.variants.llamalike_common) 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.experimental.v2.toolshed.lora) (in module penzai.toolshed.lora) LossFunction (class in penzai.experimental.v2.toolshed.basic_training) (class in penzai.toolshed.basic_training) LowRankAdapter (class in penzai.experimental.v2.toolshed.lora) (class in penzai.toolshed.lora) M make_parameter() (in module penzai.experimental.v2.nn.parameters) 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.experimental.v2.nn.dropout) (in module penzai.nn.dropout) mean() (penzai.core.named_axes.NamedArrayBase method) min() (penzai.core.named_axes.NamedArrayBase method) MisconfiguredLayerError mistral_from_huggingface_model() (in module penzai.experimental.v2.models.transformer.variants.mistral) MLP (class in penzai.example_models.simple_mlp) (class in penzai.experimental.v2.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.experimental.v2 penzai.experimental.v2.core penzai.experimental.v2.core.auto_order_types penzai.experimental.v2.core.random_stream penzai.experimental.v2.core.variables penzai.experimental.v2.models penzai.experimental.v2.models.simple_mlp penzai.experimental.v2.models.transformer penzai.experimental.v2.models.transformer.model_parts penzai.experimental.v2.models.transformer.sampling_mode penzai.experimental.v2.models.transformer.simple_decoding_loop penzai.experimental.v2.models.transformer.variants penzai.experimental.v2.models.transformer.variants.gemma penzai.experimental.v2.models.transformer.variants.gpt_neox penzai.experimental.v2.models.transformer.variants.llama penzai.experimental.v2.models.transformer.variants.llamalike_common penzai.experimental.v2.models.transformer.variants.mistral penzai.experimental.v2.nn penzai.experimental.v2.nn.attention penzai.experimental.v2.nn.basic_ops penzai.experimental.v2.nn.combinators penzai.experimental.v2.nn.dropout penzai.experimental.v2.nn.embeddings penzai.experimental.v2.nn.grouping penzai.experimental.v2.nn.layer penzai.experimental.v2.nn.layer_stack penzai.experimental.v2.nn.linear_and_affine penzai.experimental.v2.nn.parameters penzai.experimental.v2.nn.standardization penzai.experimental.v2.toolshed penzai.experimental.v2.toolshed.basic_training penzai.experimental.v2.toolshed.isolate_submodel penzai.experimental.v2.toolshed.jit_wrapper penzai.experimental.v2.toolshed.lora penzai.experimental.v2.toolshed.model_rewiring penzai.experimental.v2.toolshed.save_intermediates penzai.experimental.v2.toolshed.sharding_util penzai.experimental.v2.toolshed.unflaxify 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.experimental.v2.nn.linear_and_affine) (class in penzai.nn.linear_and_affine) NamedGroup (class in penzai.experimental.v2.nn.grouping) (class in penzai.nn.grouping) next_key() (penzai.core.random_stream.RandomStream method) (penzai.data_effects.random.RandomEffect method) (penzai.experimental.v2.core.random_stream.RandomStream 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() (in module penzai.core.named_axes) (penzai.core.named_axes.NamedArrayBase method) output_axes (penzai.experimental.v2.nn.linear_and_affine.Linear property) (penzai.nn.linear_and_affine.Linear property) output_structure() (penzai.core.layer.Layer method) P parallel_axes (penzai.experimental.v2.nn.linear_and_affine.Linear property) (penzai.nn.linear_and_affine.Linear property) Parameter (class in penzai.experimental.v2.core.variables) (class in penzai.nn.parameters) ParameterLike (class in penzai.experimental.v2.nn.parameters) (class in penzai.nn.parameters) ParameterSlot (class in penzai.experimental.v2.core.variables) ParameterValue (class in penzai.experimental.v2.core.variables) partition() (penzai.core.selectors.Selection method) patch_ipdb() (in module penzai.toolshed.patch_ipdb) path_matrix() (penzai.experimental.v2.toolshed.model_rewiring.RewireComputationPaths method) (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.experimental.v2 module penzai.experimental.v2.core module penzai.experimental.v2.core.auto_order_types module penzai.experimental.v2.core.random_stream module penzai.experimental.v2.core.variables module penzai.experimental.v2.models module penzai.experimental.v2.models.simple_mlp module penzai.experimental.v2.models.transformer module penzai.experimental.v2.models.transformer.model_parts module penzai.experimental.v2.models.transformer.sampling_mode module penzai.experimental.v2.models.transformer.simple_decoding_loop module penzai.experimental.v2.models.transformer.variants module penzai.experimental.v2.models.transformer.variants.gemma module penzai.experimental.v2.models.transformer.variants.gpt_neox module penzai.experimental.v2.models.transformer.variants.llama module penzai.experimental.v2.models.transformer.variants.llamalike_common module penzai.experimental.v2.models.transformer.variants.mistral module penzai.experimental.v2.nn module penzai.experimental.v2.nn.attention module penzai.experimental.v2.nn.basic_ops module penzai.experimental.v2.nn.combinators module penzai.experimental.v2.nn.dropout module penzai.experimental.v2.nn.embeddings module penzai.experimental.v2.nn.grouping module penzai.experimental.v2.nn.layer module penzai.experimental.v2.nn.layer_stack module penzai.experimental.v2.nn.linear_and_affine module penzai.experimental.v2.nn.parameters module penzai.experimental.v2.nn.standardization module penzai.experimental.v2.toolshed module penzai.experimental.v2.toolshed.basic_training module penzai.experimental.v2.toolshed.isolate_submodel module penzai.experimental.v2.toolshed.jit_wrapper module penzai.experimental.v2.toolshed.lora module penzai.experimental.v2.toolshed.model_rewiring module penzai.experimental.v2.toolshed.save_intermediates module penzai.experimental.v2.toolshed.sharding_util module penzai.experimental.v2.toolshed.unflaxify 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) (class in penzai.experimental.v2.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.experimental.v2.nn.linear_and_affine) (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.experimental.v2.nn.combinators) (class in penzai.nn.combinators) RewireComputationPaths (class in penzai.experimental.v2.toolshed.model_rewiring) (class in penzai.toolshed.model_rewiring) RMSLayerNorm (class in penzai.experimental.v2.nn.standardization) (class in penzai.nn.standardization) RMSStandardize (class in penzai.experimental.v2.nn.standardization) (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) SaveIntermediate (class in penzai.experimental.v2.toolshed.save_intermediates) SaveIntermediateShape (class in penzai.experimental.v2.toolshed.save_intermediates) saving_all_intermediates() (in module penzai.experimental.v2.toolshed.save_intermediates) scan() (in module penzai.core.named_axes) scoped_auto_state_var_labels() (in module penzai.experimental.v2.core.variables) ScopedStateVarLabel (class in penzai.experimental.v2.core.variables) 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.experimental.v2.nn.grouping) (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) set_value() (penzai.experimental.v2.core.variables.LabeledVariable method) sharded_init() (in module penzai.experimental.v2.toolshed.sharding_util) 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.experimental.v2.nn.basic_ops) (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.experimental.v2.nn.standardization) (class in penzai.nn.standardization) StatefulTrainer (class in penzai.experimental.v2.toolshed.basic_training) stateless_call() (penzai.experimental.v2.nn.layer.Layer method) StateVariable (class in penzai.experimental.v2.core.variables) StateVariableSlot (class in penzai.experimental.v2.core.variables) StateVariableValue (class in penzai.experimental.v2.core.variables) Static (class in penzai.toolshed.annotate_shapes) std() (penzai.core.named_axes.NamedArrayBase method) step() (penzai.experimental.v2.toolshed.basic_training.StatefulTrainer method) StochasticDropout (class in penzai.experimental.v2.nn.dropout) (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) (in module penzai.experimental.v2.models.transformer.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) TransformerBlock (class in penzai.experimental.v2.models.transformer.model_parts) TransformerFeedForward (class in penzai.experimental.v2.models.transformer.model_parts) TransformerLM (class in penzai.experimental.v2.models.transformer.model_parts) TransformerMetadata (class in penzai.experimental.v2.models.transformer.model_parts) 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) (penzai.experimental.v2.core.variables.LabeledVariable method) TreescopeRenderingFigure (class in penzai.treescope.figures) U unbind_params() (in module penzai.experimental.v2.core.variables) unbind_state_vars() (in module penzai.experimental.v2.core.variables) unbind_variables() (in module penzai.experimental.v2.core.variables) UnboundVariableError unchecked_layer_call() (in module penzai.core.layer) unflaxify_apply() (in module penzai.experimental.v2.toolshed.unflaxify) (in module penzai.toolshed.unflaxify) unfreeze_as_copy() (penzai.experimental.v2.core.variables.AbstractVariableValue method) 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) update() (penzai.experimental.v2.core.variables.AbstractVariable method) (penzai.experimental.v2.core.variables.LabeledVariable method) use_autovisualizer_if_present() (in module penzai.treescope.autovisualize) using_expansion_strategy() (in module penzai.treescope.default_renderer) V value (penzai.experimental.v2.core.variables.ParameterSlot property) (penzai.experimental.v2.core.variables.StateVariableSlot property) (penzai.experimental.v2.nn.parameters.ParameterLike property) (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) variable_jit() (in module penzai.experimental.v2.core.variables) VariableConflictError variance_scaling_initializer() (in module penzai.experimental.v2.nn.linear_and_affine) (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.experimental.v2.nn.linear_and_affine) (in module penzai.nn.linear_and_affine) zeros() (in module penzai.core.named_axes)