Index _ | A | B | C | D | E | F | G | H | I | J | K | L | M | N | O | P | Q | 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) __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) (penzai.treescope.handlers.hardcoded_structure_handlers.HardcodedStructureHandler method) (penzai.treescope.renderer.TreescopeCustomWrapperHook method) (penzai.treescope.renderer.TreescopeNodeHandler method) (penzai.treescope.renderer.TreescopeSubtreeRenderer 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 AbbreviationColor (class in penzai.treescope.foldable_representation.common_styles) 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_alias() (in module penzai.treescope.canonical_aliases) 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) append_repr_html_when_present() (in module penzai.treescope.handlers.repr_html_postprocessor) 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) ArraySpecStyle (class in penzai.treescope.handlers.penzai.shapecheck_handlers) ArrayVariableStyle (class in penzai.treescope.handlers.penzai.shapecheck_handlers) 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_expanded_part() (penzai.treescope.foldable_representation.foldable_impl.FoldableTreeNodeImpl method) (penzai.treescope.foldable_representation.part_interface.FoldableTreeNode 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 BaseBoxWithOutline (class in penzai.treescope.foldable_representation.basic_parts) BaseContentlessLeaf (class in penzai.treescope.foldable_representation.basic_parts) BaseSpanGroup (class in penzai.treescope.foldable_representation.basic_parts) BaseTaggedGroup (class in penzai.treescope.foldable_representation.basic_parts) 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() (penzai.treescope.foldable_representation.basic_parts.IndentedChildren class method) (penzai.treescope.foldable_representation.basic_parts.OnSeparateLines class method) (penzai.treescope.foldable_representation.basic_parts.Siblings class method) build_copy_button() (in module penzai.treescope.foldable_representation.common_structures) build_custom_foldable_tree_node() (in module penzai.treescope.foldable_representation.common_structures) build_field_children() (in module penzai.treescope.handlers.builtin_structure_handler) build_foldable_representation() (in module penzai.treescope.default_renderer) build_foldable_tree_node_from_children() (in module penzai.treescope.foldable_representation.common_structures) build_full_line_with_annotations() (in module penzai.treescope.foldable_representation.basic_parts) build_one_line_tree_node() (in module penzai.treescope.foldable_representation.common_structures) 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) CanonicalAliasEnvironment (class in penzai.treescope.canonical_aliases) canonicalize() (penzai.core.named_axes.NamedArrayBase method) capture() (penzai.toolshed.unflaxify.ArgsAndKwargs class method) CastToDType (class in penzai.nn.basic_ops) check_for_shared_values() (in module penzai.treescope.handlers.shared_value_postprocessor) 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) collapsed_width (penzai.treescope.foldable_representation.part_interface.RenderableTreePart property) collecting_deferred_renderings() (in module penzai.treescope.foldable_representation.foldable_impl) CollectingSideOutputs (class in penzai.data_effects.side_output) color_from_string() (in module penzai.core.formatting_util) ColoredBorderIndentedChildren (class in penzai.treescope.foldable_representation.common_styles) ColoredBottomLineSpanGroup (class in penzai.treescope.foldable_representation.common_styles) ColoredSingleLineSpanGroup (class in penzai.treescope.foldable_representation.common_styles) ColoredTopLineSpanGroup (class in penzai.treescope.foldable_representation.common_styles) combine() (in module penzai.core.partitioning) CommentColor (class in penzai.treescope.foldable_representation.common_styles) CommentColorWhenExpanded (class in penzai.treescope.foldable_representation.common_styles) compress() (penzai.core.named_axes.NamedArrayBase method) compress_html() (in module penzai.treescope.html_compression) compute_training_outputs_and_updates() (in module penzai.toolshed.basic_training) compute_truncated_shape() (in module penzai.treescope.ndarray_summarization) 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) CSSStyleRule (class in penzai.treescope.foldable_representation.part_interface) cumprod() (penzai.core.named_axes.NamedArrayBase method) cumsum() (penzai.core.named_axes.NamedArrayBase method) CustomTextColor (class in penzai.treescope.foldable_representation.common_styles) CustomTreescopeVisualization (class in penzai.treescope.autovisualize) D DashedGrayOutlineBox (class in penzai.treescope.foldable_representation.common_styles) dataclass_from_attributes() (in module penzai.core.dataclass_util) decompression_preamble() (in module penzai.treescope.html_compression) 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) default_well_known_filter() (in module penzai.treescope.canonical_aliases) DeferredPlaceholder (class in penzai.treescope.foldable_representation.foldable_impl) DeferredPlaceholderStyle (class in penzai.treescope.foldable_representation.common_styles) DeferredWithThunk (class in penzai.treescope.foldable_representation.foldable_impl) DeferringToChild (class in penzai.treescope.foldable_representation.basic_parts) 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) display_selection_streaming() (in module penzai.treescope.selection_rendering) display_streaming_as_root() (in module penzai.treescope.foldable_representation.foldable_impl) 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) DynamicSharedCheck (class in penzai.treescope.handlers.shared_value_postprocessor) 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) EmbeddedIFrame (class in penzai.treescope.foldable_representation.embedded_iframe) EmbeddingDecode (class in penzai.nn.embeddings) EmbeddingLookup (class in penzai.nn.embeddings) EmbeddingTable (class in penzai.nn.embeddings) EmptyPart (class in penzai.treescope.foldable_representation.basic_parts) EmptyWithHeightGuess (class in penzai.treescope.foldable_representation.common_structures) enable() (penzai.nn.dropout.DisabledDropout method) enable_interactive_context() (in module penzai.core.context) ErrorColor (class in penzai.treescope.foldable_representation.common_styles) escape_html_attribute() (in module penzai.treescope.html_escaping) expand_for_balanced_layout() (in module penzai.treescope.foldable_representation.layout_algorithms) expand_to_depth() (in module penzai.treescope.foldable_representation.layout_algorithms) expand_to_tags() (in module penzai.treescope.foldable_representation.layout_algorithms) ExpandState (class in penzai.treescope.foldable_representation.part_interface) extended_with() (penzai.treescope.renderer.TreescopeRenderer method) F fake_placeholder_foldable() (in module penzai.treescope.foldable_representation.common_structures) faster_array_repr() (in module penzai.treescope.ndarray_summarization) flatten() (penzai.core.named_axes.NamedArrayBase method) flatten_selected_selections() (penzai.core.selectors.Selection method) floating_named() (penzai.core.shapecheck.ArraySpec class method) foldables_in_this_part() (penzai.treescope.foldable_representation.part_interface.FoldableTreeNode method) (penzai.treescope.foldable_representation.part_interface.RenderableTreePart method) FoldableTreeNode (class in penzai.treescope.foldable_representation.part_interface) FoldableTreeNodeImpl (class in penzai.treescope.foldable_representation.foldable_impl) FoldCondition (class in penzai.treescope.foldable_representation.basic_parts) for_tokenizer() (penzai.treescope.arrayviz.array_autovisualizer.ArrayAutovisualizer class method) format_source_location() (in module penzai.treescope.handlers.function_reflection_handlers) 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_dtype_name() (in module penzai.treescope.ndarray_summarization) get_effect_color() (in module penzai.data_effects.effect_base) get_expand_state() (penzai.treescope.foldable_representation.foldable_impl.FoldableTreeNodeImpl method) (penzai.treescope.foldable_representation.part_interface.FoldableTreeNode method) get_keypaths() (penzai.core.selectors.Selection method) get_sequence() (penzai.core.selectors.Selection method) GetPieceSize() (penzai.toolshed.token_visualization.SentencePieceProcessorLike method) H handle_anything_with_repr() (in module penzai.treescope.handlers.generic_repr_handler) handle_arbitrary_pytrees() (in module penzai.treescope.handlers.generic_pytree_handler) handle_arraystructures() (in module penzai.treescope.handlers.penzai.shapecheck_handlers) handle_builtin_atoms() (in module penzai.treescope.handlers.builtin_atom_handler) handle_builtin_structures() (in module penzai.treescope.handlers.builtin_structure_handler) handle_code_objects_with_reflection() (in module penzai.treescope.handlers.function_reflection_handlers) handle_data_effects_objects() (in module penzai.treescope.handlers.penzai.data_effects_handlers) handle_dtype_instances() (in module penzai.treescope.handlers.ndarray_handler) handle_layers() (in module penzai.treescope.handlers.penzai.layer_handler) handle_local_states() (in module penzai.data_effects.local_state) handle_named_arrays() (in module penzai.treescope.handlers.penzai.named_axes_handlers) handle_ndarrays() (in module penzai.treescope.handlers.ndarray_handler) handle_structs() (in module penzai.treescope.handlers.penzai.struct_handler) 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) HardcodedStructureHandler (class in penzai.treescope.handlers.hardcoded_structure_handlers) HasExplicitFields (class in penzai.treescope.handlers.hardcoded_structure_handlers) HasFieldsInClassAttr (class in penzai.treescope.handlers.hardcoded_structure_handlers) HasFieldsLikeInit (class in penzai.treescope.handlers.hardcoded_structure_handlers) HasReprHtml (class in penzai.treescope.foldable_representation.embedded_iframe) heuristic_strip_javascript_comments() (in module penzai.treescope.html_escaping) 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) html_setup_parts() (penzai.treescope.foldable_representation.part_interface.RenderableTreePart method) HtmlContextForSetup (class in penzai.treescope.foldable_representation.part_interface) HyperlinkTarget (class in penzai.treescope.foldable_representation.foldable_impl) 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) IndentedChildren (class in penzai.treescope.foldable_representation.basic_parts) infer_balanced_truncation() (in module penzai.treescope.ndarray_summarization) 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) is_weak() (penzai.treescope.foldable_representation.part_interface.ExpandState method) IsControl() (penzai.toolshed.token_visualization.SentencePieceProcessorLike method) IsEnumLike (class in penzai.treescope.handlers.hardcoded_structure_handlers) IsolatedSubmodel (class in penzai.toolshed.isolate_submodel) IsolationCapturer (class in penzai.toolshed.isolate_submodel) item() (penzai.core.named_axes.NamedArrayBase method) J JavaScriptDefn (class in penzai.treescope.foldable_representation.part_interface) Jitted (class in penzai.toolshed.jit_wrapper) K key_for_field() (penzai.core.struct.Struct method) KeywordColor (class in penzai.treescope.handlers.builtin_atom_handler) 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) LocalNamePath (class in penzai.treescope.canonical_aliases) LocalStateEffect (class in penzai.data_effects.local_state) LocalStateEffectImpl (class in penzai.data_effects.local_state) lookup_alias() (in module penzai.treescope.canonical_aliases) 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_defer_rendering() (in module penzai.treescope.foldable_representation.foldable_impl) maybe_dropout() (in module penzai.nn.dropout) maybe_local_module_name() (in module penzai.treescope.canonical_aliases) maybe_qualified_type_name() (in module penzai.treescope.foldable_representation.common_structures) 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.canonical_aliases penzai.treescope.copypaste_fallback penzai.treescope.default_renderer penzai.treescope.figures penzai.treescope.foldable_representation penzai.treescope.foldable_representation.basic_parts penzai.treescope.foldable_representation.common_structures penzai.treescope.foldable_representation.common_styles penzai.treescope.foldable_representation.embedded_iframe penzai.treescope.foldable_representation.foldable_impl penzai.treescope.foldable_representation.layout_algorithms penzai.treescope.foldable_representation.part_interface penzai.treescope.handlers penzai.treescope.handlers.builtin_atom_handler penzai.treescope.handlers.builtin_structure_handler penzai.treescope.handlers.canonical_alias_postprocessor penzai.treescope.handlers.function_reflection_handlers penzai.treescope.handlers.generic_pytree_handler penzai.treescope.handlers.generic_repr_handler penzai.treescope.handlers.hardcoded_structure_handlers penzai.treescope.handlers.ndarray_handler penzai.treescope.handlers.penzai penzai.treescope.handlers.penzai.data_effects_handlers penzai.treescope.handlers.penzai.layer_handler penzai.treescope.handlers.penzai.named_axes_handlers penzai.treescope.handlers.penzai.shapecheck_handlers penzai.treescope.handlers.penzai.struct_handler penzai.treescope.handlers.repr_html_postprocessor penzai.treescope.handlers.shared_value_postprocessor penzai.treescope.html_compression penzai.treescope.html_escaping penzai.treescope.ndarray_summarization penzai.treescope.renderer penzai.treescope.selection_rendering penzai.treescope.treescope_ipython ModuleAttributePath (class in penzai.treescope.canonical_aliases) 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_array_and_contained_type_summary() (in module penzai.treescope.handlers.penzai.named_axes_handlers) 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) newlines_in_expanded_parent (penzai.treescope.foldable_representation.part_interface.RenderableTreePart property) next_key() (penzai.core.random_stream.RandomStream method) (penzai.data_effects.random.RandomEffect method) nmap() (in module penzai.core.named_axes) NodeHyperlink (class in penzai.treescope.foldable_representation.foldable_impl) nonzero() (penzai.core.named_axes.NamedArrayBase method) NotInThisPartition (class in penzai.core.partitioning) NotRoundtrippable (class in penzai.treescope.copypaste_fallback) NumberColor (class in penzai.treescope.handlers.builtin_atom_handler) O oklch_color() (in module penzai.core.formatting_util) ones() (in module penzai.core.named_axes) OnSeparateLines (class in penzai.treescope.foldable_representation.basic_parts) 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) parse_color_and_pattern() (in module penzai.treescope.handlers.builtin_structure_handler) partition() (penzai.core.selectors.Selection method) patch_ipdb() (in module penzai.toolshed.patch_ipdb) path_matrix() (penzai.toolshed.model_rewiring.RewireComputationPaths method) PatternedBottomLineSpanGroup (class in penzai.treescope.foldable_representation.common_styles) PatternedSingleLineSpanGroup (class in penzai.treescope.foldable_representation.common_styles) PatternedTopLineSpanGroup (class in penzai.treescope.foldable_representation.common_styles) 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.canonical_aliases module penzai.treescope.copypaste_fallback module penzai.treescope.default_renderer module penzai.treescope.figures module penzai.treescope.foldable_representation module penzai.treescope.foldable_representation.basic_parts module penzai.treescope.foldable_representation.common_structures module penzai.treescope.foldable_representation.common_styles module penzai.treescope.foldable_representation.embedded_iframe module penzai.treescope.foldable_representation.foldable_impl module penzai.treescope.foldable_representation.layout_algorithms module penzai.treescope.foldable_representation.part_interface module penzai.treescope.handlers module penzai.treescope.handlers.builtin_atom_handler module penzai.treescope.handlers.builtin_structure_handler module penzai.treescope.handlers.canonical_alias_postprocessor module penzai.treescope.handlers.function_reflection_handlers module penzai.treescope.handlers.generic_pytree_handler module penzai.treescope.handlers.generic_repr_handler module penzai.treescope.handlers.hardcoded_structure_handlers module penzai.treescope.handlers.ndarray_handler module penzai.treescope.handlers.penzai module penzai.treescope.handlers.penzai.data_effects_handlers module penzai.treescope.handlers.penzai.layer_handler module penzai.treescope.handlers.penzai.named_axes_handlers module penzai.treescope.handlers.penzai.shapecheck_handlers module penzai.treescope.handlers.penzai.struct_handler module penzai.treescope.handlers.repr_html_postprocessor module penzai.treescope.handlers.shared_value_postprocessor module penzai.treescope.html_compression module penzai.treescope.html_escaping module penzai.treescope.ndarray_summarization module penzai.treescope.renderer module penzai.treescope.selection_rendering module penzai.treescope.treescope_ipython module pick_nth_selected() (penzai.core.selectors.Selection method) populate_from_public_api() (in module penzai.treescope.canonical_aliases) positional_shape (penzai.core.named_axes.NamedArrayBase property) prefill() (in module penzai.example_models.gemma.simple_decoding_loop) prefix_filter() (in module penzai.treescope.canonical_aliases) 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 PyTreeNodeFieldName (class in penzai.treescope.handlers.penzai.struct_handler) Q QualifiedTypeNameSpanGroup (class in penzai.treescope.foldable_representation.common_styles) 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) relative_alias_names() (in module penzai.treescope.canonical_aliases) 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_dataclass_constructor() (in module penzai.treescope.handlers.builtin_structure_handler) render_not_roundtrippable() (in module penzai.treescope.copypaste_fallback) render_selection_to_foldable_representation() (in module penzai.treescope.selection_rendering) render_sharded_shape() (in module penzai.treescope.arrayviz.arrayviz) render_short_struct_summary() (in module penzai.treescope.handlers.penzai.struct_handler) render_struct_constructor() (in module penzai.treescope.handlers.penzai.struct_handler) render_to_html() (in module penzai.treescope.default_renderer) (penzai.treescope.foldable_representation.part_interface.RenderableTreePart method) render_to_html_as_root() (in module penzai.treescope.foldable_representation.foldable_impl) render_to_text() (in module penzai.treescope.default_renderer) (penzai.treescope.foldable_representation.part_interface.RenderableTreePart method) render_to_text_as_root() (in module penzai.treescope.foldable_representation.foldable_impl) RenderableAndLineAnnotations (class in penzai.treescope.foldable_representation.part_interface) RenderableTreePart (class in penzai.treescope.foldable_representation.part_interface) RendersAsRootInIPython (class in penzai.treescope.figures) repeat() (penzai.core.named_axes.NamedArrayBase method) replace_with_canonical_aliases() (in module penzai.treescope.handlers.canonical_alias_postprocessor) reshape() (penzai.core.named_axes.NamedArrayBase method) Residual (class in penzai.nn.combinators) retrieve() (penzai.treescope.canonical_aliases.LocalNamePath method) (penzai.treescope.canonical_aliases.ModuleAttributePath method) 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) RoundtripCondition (class in penzai.treescope.foldable_representation.basic_parts) run_and_interleave_intermediates() (in module penzai.toolshed.interleave_intermediates) S safe_to_summarize() (in module penzai.treescope.ndarray_summarization) SamplingState (class in penzai.example_models.gemma.simple_decoding_loop) ScopedSelectableAnnotation (class in penzai.treescope.foldable_representation.basic_parts) 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) SelectedNodeTracker (class in penzai.treescope.selection_rendering) Selection (class in penzai.core.selectors) SelectionBoundaryBox (class in penzai.treescope.selection_rendering) SelectionBoundaryLabel (class in penzai.treescope.selection_rendering) SelectionBoundaryTag (class in penzai.treescope.selection_rendering) SelectionHole (class in penzai.core.selectors) SelectionQuote (class in penzai.core.selectors) SelectionTaggedGroup (class in penzai.treescope.selection_rendering) 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_expand_state() (penzai.treescope.foldable_representation.foldable_impl.FoldableTreeNodeImpl method) (penzai.treescope.foldable_representation.part_interface.FoldableTreeNode method) set_interactive() (penzai.core.context.ContextualValue method) set_scoped() (penzai.core.context.ContextualValue method) set_sequence() (penzai.core.selectors.Selection method) setup_shared_value_context() (in module penzai.treescope.handlers.shared_value_postprocessor) 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) SharedWarningLabel (class in penzai.treescope.handlers.shared_value_postprocessor) 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) Siblings (class in penzai.treescope.foldable_representation.basic_parts) siblings() (in module penzai.treescope.foldable_representation.basic_parts) siblings_with_annotations() (in module penzai.treescope.foldable_representation.basic_parts) 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) StringCopyButton (class in penzai.treescope.foldable_representation.foldable_impl) StringLiteralColor (class in penzai.treescope.handlers.builtin_atom_handler) Struct (class in penzai.core.struct) struct_attr_style_fn_for_fields() (in module penzai.treescope.handlers.penzai.struct_handler) StructStaticMetadata (class in penzai.core.struct) StructureMismatchError styled() (in module penzai.treescope.figures) sum() (penzai.core.named_axes.NamedArrayBase method) SummarizableCondition (class in penzai.treescope.foldable_representation.basic_parts) summarization_threshold (in module penzai.treescope.ndarray_summarization) summarize_ndarray() (in module penzai.treescope.ndarray_summarization) 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) tags_in_this_part (penzai.treescope.foldable_representation.part_interface.RenderableTreePart property) 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 (class in penzai.treescope.foldable_representation.basic_parts) text_on_color() (in module penzai.treescope.arrayviz.arrayviz) TmpPosAxisMarker (class in penzai.core.named_axes) to_foldable_representation() (penzai.treescope.renderer.TreescopeRenderer method) to_html() (in module penzai.treescope.foldable_representation.embedded_iframe) (penzai.treescope.renderer.TreescopeRenderer method) to_text() (penzai.treescope.renderer.TreescopeRenderer method) 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) TreescopeCustomWrapperHook (class in penzai.treescope.renderer) TreescopeNodeHandler (class in penzai.treescope.renderer) TreescopeRenderer (class in penzai.treescope.renderer) TreescopeRenderingFigure (class in penzai.treescope.figures) TreescopeSubtreeRenderer (class in penzai.treescope.renderer) truncate_array_and_mask() (in module penzai.treescope.ndarray_summarization) 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) update_lazy_aliases() (in module penzai.treescope.canonical_aliases) 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) VerticalSpace (class in penzai.treescope.foldable_representation.basic_parts) 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) WithBlockColor (class in penzai.treescope.foldable_representation.common_styles) WithBlockPattern (class in penzai.treescope.foldable_representation.common_styles) WithConstantSideInputs (class in penzai.data_effects.side_input) WithDynamicSharedPip (class in penzai.treescope.handlers.shared_value_postprocessor) WithFrozenRandomState (class in penzai.data_effects.random) WithFunctionalLocalState (class in penzai.data_effects.local_state) WithHoverTooltip (class in penzai.treescope.foldable_representation.basic_parts) without_repeated_whitespace() (in module penzai.treescope.html_escaping) 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)