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.struct.AbstractStructMetaclass method) (penzai.deprecated.v1.core.layer.Layer method) (penzai.deprecated.v1.example_models.gemma.model_core.GemmaTransformer method) (penzai.deprecated.v1.example_models.gemma.sampling_mode.GemmaKVCachingTransformer method) (penzai.deprecated.v1.nn.attention.ApplyAttentionMask method) (penzai.deprecated.v1.nn.attention.Attention method) (penzai.deprecated.v1.nn.attention.KVCachingAttention method) (penzai.deprecated.v1.nn.combinators.Residual method) (penzai.deprecated.v1.nn.dropout.StochasticDropout method) (penzai.deprecated.v1.nn.embeddings.EmbeddingDecode method) (penzai.deprecated.v1.nn.embeddings.EmbeddingLookup method) (penzai.deprecated.v1.nn.grouping.CheckedSequential method) (penzai.deprecated.v1.nn.grouping.CheckStructure method) (penzai.deprecated.v1.nn.grouping.Identity method) (penzai.deprecated.v1.nn.grouping.NamedGroup method) (penzai.deprecated.v1.nn.grouping.Sequential method) (penzai.deprecated.v1.nn.linear_and_affine.AddBias method) (penzai.deprecated.v1.nn.linear_and_affine.ConstantRescale method) (penzai.deprecated.v1.nn.linear_and_affine.Linear method) (penzai.deprecated.v1.nn.linear_and_affine.LinearOperatorWeightInitializer method) (penzai.deprecated.v1.nn.linear_and_affine.NamedEinsum method) (penzai.deprecated.v1.nn.standardization.RMSStandardize method) (penzai.deprecated.v1.nn.standardization.Standardize method) (penzai.deprecated.v1.toolshed.basic_training.LossFunction method) (penzai.deprecated.v1.toolshed.basic_training.TrainStepFunction method) (penzai.deprecated.v1.toolshed.unflaxify.InterceptedFlaxModuleMethod method) (penzai.models.transformer.model_parts.TransformerLM method) (penzai.models.transformer.sampling_mode.KVCachingTransformerLM method) (penzai.nn.attention.ApplyCausalAttentionMask method) (penzai.nn.attention.ApplyCausalSlidingWindowAttentionMask method) (penzai.nn.attention.ApplyExplicitAttentionMask 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.layer.Layer method) (penzai.nn.layer_stack.LayerStack 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.unflaxify.InterceptedFlaxModuleMethod method) __divmod__() (penzai.core.named_axes.NamedArrayBase method) __enter__() (penzai.deprecated.v1.core.random_stream.RandomStream method) __exit__() (penzai.deprecated.v1.core.random_stream.RandomStream method) __floordiv__() (penzai.core.named_axes.NamedArrayBase method) __ge__() (penzai.core.named_axes.NamedArrayBase method) __getattr__() (penzai.deprecated.v1.data_effects.effect_base.EffectRequest method) (penzai.deprecated.v1.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.shapecheck.StructureMismatchError method) (penzai.core.struct.PyTreeDataclassSafetyError method) (penzai.core.variables.LabeledVariable method) (penzai.core.variables.StateVariable method) (penzai.core.variables.UnboundVariableError method) (penzai.core.variables.VariableConflictError method) (penzai.deprecated.v1.core.layer.MisconfiguredLayerError method) (penzai.deprecated.v1.data_effects.effect_base.UnhandledEffectError method) (penzai.deprecated.v1.nn.parameters.UninitializedParameter method) (penzai.deprecated.v1.nn.parameters.UninitializedParameterError method) __init_subclass__() (penzai.deprecated.v1.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.shapecheck.StructureMismatchError method) (penzai.core.struct.PyTreeDataclassSafetyError method) (penzai.core.variables.UnboundVariableError method) (penzai.core.variables.VariableConflictError method) (penzai.deprecated.v1.core.layer.MisconfiguredLayerError method) (penzai.deprecated.v1.data_effects.effect_base.UnhandledEffectError method) (penzai.deprecated.v1.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) __treescope_ndarray_adapter__() (penzai.core.named_axes.NamedArrayBase method) __treescope_repr__() (penzai.core.named_axes.NamedArrayBase method) __treescope_root_repr__() (penzai.core.selectors.Selection 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.core.variables) AbstractVariableSlot (class in penzai.core.variables) AbstractVariableValue (class in penzai.core.variables) add_parameter_prefix() (in module penzai.deprecated.v1.nn.parameters) AddBias (class in penzai.deprecated.v1.nn.linear_and_affine) (class in penzai.nn.linear_and_affine) advance_one_token() (in module penzai.deprecated.v1.example_models.gemma.simple_decoding_loop) Affine (class in penzai.deprecated.v1.nn.linear_and_affine) (class in penzai.nn.linear_and_affine) all() (penzai.core.named_axes.NamedArrayBase method) all_handler_ids() (in module penzai.deprecated.v1.data_effects.effect_base) annotate_shapes() (in module penzai.deprecated.v1.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.deprecated.v1.nn.attention) ApplyCausalAttentionMask (class in penzai.nn.attention) ApplyCausalSlidingWindowAttentionMask (class in penzai.nn.attention) ApplyExplicitAttentionMask (class in penzai.nn.attention) ApplyRoPE (class in penzai.deprecated.v1.nn.embeddings) (class in penzai.nn.embeddings) ApplyRoPEToSubset (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.deprecated.v1.toolshed.unflaxify) (class in penzai.toolshed.unflaxify) argsort() (penzai.core.named_axes.NamedArrayBase method) ArraySpec (class in penzai.core.shapecheck) as_array_structure() (in module penzai.core.shapecheck) as_empty_parameter() (penzai.deprecated.v1.nn.parameters.UninitializedParameter method) as_frozen() (penzai.deprecated.v1.data_effects.local_state.InitialLocalStateRequest method) as_namedarrayview() (penzai.core.named_axes.NamedArrayBase method) ask() (penzai.deprecated.v1.data_effects.side_input.SideInputEffect method) assert_count_is() (penzai.core.selectors.Selection method) assert_no_parameter_slots() (in module penzai.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.deprecated.v1.nn.parameters) Attention (class in penzai.deprecated.v1.nn.attention) (class in penzai.nn.attention) AttentionTypeGlobalCausal (class in penzai.models.transformer.variants.llamalike_common) AttentionTypeSlidingWindowCausal (class in penzai.models.transformer.variants.llamalike_common) attributes_dict() (penzai.core.struct.Struct method) AutoNmapModuleWrapper (class in penzai.toolshed.auto_nmap) AutoOrderedAcrossTypes (class in penzai.core.auto_order_types) AutoStateVarLabel (class in penzai.core.variables) B bind_variables() (in module penzai.core.variables) (penzai.nn.layer.Layer method) BranchAndAddTogether (class in penzai.deprecated.v1.nn.combinators) (class in penzai.nn.combinators) BranchAndMultiplyTogether (class in penzai.deprecated.v1.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.deprecated.v1.data_effects.effect_base) build_gpt_neox_attention() (in module penzai.models.transformer.variants.gpt_neox) build_gpt_neox_block() (in module penzai.models.transformer.variants.gpt_neox) build_gpt_neox_feedforward() (in module penzai.models.transformer.variants.gpt_neox) build_gpt_neox_transformer() (in module penzai.models.transformer.variants.gpt_neox) build_llamalike_attention() (in module penzai.models.transformer.variants.llamalike_common) build_llamalike_block() (in module penzai.models.transformer.variants.llamalike_common) build_llamalike_feedforward() (in module penzai.models.transformer.variants.llamalike_common) build_llamalike_transformer() (in module penzai.models.transformer.variants.llamalike_common) build_train_step_fn() (in module penzai.deprecated.v1.toolshed.basic_training) C call_and_extract_submodel() (in module penzai.deprecated.v1.toolshed.isolate_submodel) (in module penzai.toolshed.isolate_submodel) CalledWithManyStructures (class in penzai.deprecated.v1.toolshed.annotate_shapes) canonicalize() (penzai.core.named_axes.NamedArrayBase method) capture() (penzai.deprecated.v1.toolshed.unflaxify.ArgsAndKwargs class method) (penzai.toolshed.unflaxify.ArgsAndKwargs class method) CastToDType (class in penzai.deprecated.v1.nn.basic_ops) (class in penzai.nn.basic_ops) check_layer() (in module penzai.deprecated.v1.toolshed.check_layers_by_tracing) check_no_duplicated_parameters() (in module penzai.deprecated.v1.nn.parameters) check_structure() (in module penzai.core.shapecheck) check_unique_axis_names_for_initializer() (in module penzai.deprecated.v1.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.deprecated.v1.core.layer) CheckedSequential (class in penzai.deprecated.v1.nn.grouping) (class in penzai.nn.grouping) Checkpointed (class in penzai.toolshed.gradient_checkpointing) CheckStructure (class in penzai.deprecated.v1.nn.grouping) (class in penzai.nn.grouping) choose() (penzai.core.named_axes.NamedArrayBase method) clip() (penzai.core.named_axes.NamedArrayBase method) CollectingSideOutputs (class in penzai.deprecated.v1.data_effects.side_output) combine() (in module penzai.core.partitioning) compress() (penzai.core.named_axes.NamedArrayBase method) compute_training_outputs_and_updates() (in module penzai.deprecated.v1.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.deprecated.v1.nn.linear_and_affine) (in module penzai.nn.linear_and_affine) ConstantRescale (class in penzai.deprecated.v1.nn.linear_and_affine) (class in penzai.nn.linear_and_affine) ConstrainSharding (class in penzai.deprecated.v1.toolshed.sharding_util) (class in penzai.toolshed.sharding_util) ConstrainShardingByName (class in penzai.deprecated.v1.toolshed.sharding_util) (class in penzai.toolshed.sharding_util) contract() (in module penzai.deprecated.v1.nn.linear_and_affine) (in module penzai.nn.linear_and_affine) count() (penzai.core.selectors.Selection method) cumprod() (penzai.core.named_axes.NamedArrayBase method) cumsum() (penzai.core.named_axes.NamedArrayBase method) D derive_param_key() (in module penzai.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.deprecated.v1.nn.dropout.StochasticDropout method) (penzai.nn.dropout.StochasticDropout method) DisabledDropout (class in penzai.deprecated.v1.nn.dropout) (class in penzai.nn.dropout) displayhook_with_ipython() (in module penzai.toolshed.patch_ipdb) dot() (penzai.core.named_axes.NamedArrayBase method) DropoutMLP (class in penzai.deprecated.v1.example_models.simple_mlp) (class in penzai.models.simple_mlp) dtype (penzai.core.named_axes.NamedArrayBase property) (penzai.core.shapecheck.ArraySpec attribute) E effect_protocol() (penzai.deprecated.v1.data_effects.effect_base.EffectHandler class method) (penzai.deprecated.v1.data_effects.effect_base.EffectRequest class method) (penzai.deprecated.v1.data_effects.effect_base.EffectRuntimeImpl class method) (penzai.deprecated.v1.data_effects.effect_base.HandledEffectRef class method) EffectHandler (class in penzai.deprecated.v1.data_effects.effect_base) EffectRequest (class in penzai.deprecated.v1.data_effects.effect_base) EffectRuntimeImpl (class in penzai.deprecated.v1.data_effects.effect_base) Elementwise (class in penzai.deprecated.v1.nn.basic_ops) (class in penzai.nn.basic_ops) embed_shared_state_requests() (in module penzai.deprecated.v1.data_effects.local_state) EmbeddingDecode (class in penzai.deprecated.v1.nn.embeddings) (class in penzai.nn.embeddings) EmbeddingLookup (class in penzai.deprecated.v1.nn.embeddings) (class in penzai.nn.embeddings) EmbeddingTable (class in penzai.deprecated.v1.nn.embeddings) (class in penzai.nn.embeddings) enable() (penzai.deprecated.v1.nn.dropout.DisabledDropout method) (penzai.nn.dropout.DisabledDropout method) F flatten() (penzai.core.named_axes.NamedArrayBase method) flatten_selected_selections() (penzai.core.selectors.Selection method) FlaxVarLabel (class in penzai.toolshed.unflaxify) floating_named() (penzai.core.shapecheck.ArraySpec class method) free_effect_types() (in module penzai.deprecated.v1.data_effects.effect_base) freeze() (penzai.core.variables.AbstractVariable method) freeze_local_states() (in module penzai.deprecated.v1.data_effects.local_state) freeze_params() (in module penzai.core.variables) freeze_state_vars() (in module penzai.core.variables) freeze_variables() (in module penzai.core.variables) From (class in penzai.deprecated.v1.toolshed.model_rewiring) (class in penzai.toolshed.model_rewiring) from_attributes() (penzai.core.struct.Struct class method) from_base_key() (penzai.core.random_stream.RandomStream class method) from_basic_segments() (penzai.deprecated.v1.example_models.gemma.model_core.GemmaInputs class method) from_basic_subsegments() (penzai.deprecated.v1.example_models.gemma.sampling_mode.GemmaKVCachingInputs class method) from_config() (penzai.deprecated.v1.data_effects.side_output.TellIntermediate class method) (penzai.deprecated.v1.example_models.gemma.model_core.GemmaAttention class method) (penzai.deprecated.v1.example_models.gemma.model_core.GemmaFeedForward class method) (penzai.deprecated.v1.example_models.gemma.model_core.GemmaTransformer class method) (penzai.deprecated.v1.example_models.gemma.model_core.GemmaTransformerBlock class method) (penzai.deprecated.v1.nn.attention.ApplyAttentionMask class method) (penzai.deprecated.v1.nn.embeddings.ApplyRoPE class method) (penzai.deprecated.v1.nn.embeddings.EmbeddingTable class method) (penzai.deprecated.v1.nn.linear_and_affine.AddBias class method) (penzai.deprecated.v1.nn.linear_and_affine.Linear class method) (penzai.deprecated.v1.nn.standardization.LayerNorm class method) (penzai.deprecated.v1.nn.standardization.RMSLayerNorm 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.deprecated.v1.toolshed.lora.LowRankAdapter class method) (penzai.toolshed.lora.LowRankAdapter class method) from_pretrained() (penzai.deprecated.v1.example_models.gemma.model_core.GemmaTransformer class method) from_sublayer_builder() (penzai.nn.layer_stack.LayerStack class method) from_uncached() (penzai.deprecated.v1.example_models.gemma.sampling_mode.GemmaKVCachingAttention class method) (penzai.deprecated.v1.example_models.gemma.sampling_mode.GemmaKVCachingTransformer class method) (penzai.deprecated.v1.nn.attention.KVCachingAttention class method) (penzai.models.transformer.sampling_mode.KVCachingTransformerLM class method) (penzai.nn.attention.KVCachingAttention class method) from_uninitialized() (penzai.deprecated.v1.nn.parameters.ShareableUninitializedParameter class method) FrozenLocalStateRequest (class in penzai.deprecated.v1.data_effects.local_state) FrozenParameter (class in penzai.deprecated.v1.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.models.transformer.variants.gemma) GemmaAttention (class in penzai.deprecated.v1.example_models.gemma.model_core) GemmaFeedForward (class in penzai.deprecated.v1.example_models.gemma.model_core) GemmaInputs (class in penzai.deprecated.v1.example_models.gemma.model_core) GemmaKVCachingAttention (class in penzai.deprecated.v1.example_models.gemma.sampling_mode) GemmaKVCachingInputs (class in penzai.deprecated.v1.example_models.gemma.sampling_mode) GemmaKVCachingState (class in penzai.deprecated.v1.example_models.gemma.sampling_mode) GemmaKVCachingTransformer (class in penzai.deprecated.v1.example_models.gemma.sampling_mode) GemmaTransformer (class in penzai.deprecated.v1.example_models.gemma.model_core) GemmaTransformerBlock (class in penzai.deprecated.v1.example_models.gemma.model_core) GemmaTransformerConfig (class in penzai.deprecated.v1.example_models.gemma.model_core) get() (penzai.core.selectors.Selection method) (penzai.deprecated.v1.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.deprecated.v1.data_effects.effect_base) get_keypaths() (penzai.core.selectors.Selection method) get_sequence() (penzai.core.selectors.Selection method) get_slot() (penzai.core.variables.AbstractVariable method) (penzai.core.variables.AbstractVariableValue method) GetPieceSize() (penzai.toolshed.token_visualization.SentencePieceProcessorLike method) gpt_neox_from_huggingface_model() (in module penzai.models.transformer.variants.gpt_neox) GPTNeoXTransformerConfig (class in penzai.models.transformer.variants.gpt_neox) H handle_local_states() (in module penzai.deprecated.v1.data_effects.local_state) HandledEffectRef (class in penzai.deprecated.v1.data_effects.effect_base) HandledLocalStateRef (class in penzai.deprecated.v1.data_effects.local_state) HandledRandomRef (class in penzai.deprecated.v1.data_effects.random) HandledSideInputRef (class in penzai.deprecated.v1.data_effects.side_input) HandledSideOutputRef (class in penzai.deprecated.v1.data_effects.side_output) handler_id() (penzai.deprecated.v1.data_effects.effect_base.EffectRuntimeImpl method) handling() (penzai.deprecated.v1.data_effects.random.WithFrozenRandomState class method) (penzai.deprecated.v1.data_effects.random.WithRandomKeyFromArg class method) (penzai.deprecated.v1.data_effects.random.WithStatefulRandomKey class method) (penzai.deprecated.v1.data_effects.side_input.WithConstantSideInputs class method) (penzai.deprecated.v1.data_effects.side_input.WithSideInputsFromInputTuple class method) (penzai.deprecated.v1.data_effects.side_output.CollectingSideOutputs class method) hoist_constant_side_inputs() (in module penzai.deprecated.v1.data_effects.side_input) hoist_shared_state_requests() (in module penzai.deprecated.v1.data_effects.local_state) HoistedTag (class in penzai.deprecated.v1.data_effects.side_input) I Identity (class in penzai.deprecated.v1.nn.grouping) (class in penzai.nn.grouping) IdentityWithSavedActivations (class in penzai.deprecated.v1.toolshed.interleave_intermediates) IdToPiece() (penzai.toolshed.token_visualization.SentencePieceProcessorLike method) imag (penzai.core.named_axes.NamedArrayBase attribute) infer_or_check_handler_id() (in module penzai.deprecated.v1.data_effects.effect_base) initial_state() (penzai.deprecated.v1.toolshed.basic_training.TrainState class method) initialize() (penzai.deprecated.v1.nn.parameters.UninitializedParameter method) initialize_parameters() (in module penzai.deprecated.v1.nn.parameters) initialize_parameters_sharded() (in module penzai.deprecated.v1.toolshed.sharding_util) initialize_with_value() (penzai.deprecated.v1.nn.parameters.UninitializedParameter method) InitialLocalStateRequest (class in penzai.deprecated.v1.data_effects.local_state) inline_anonymous_sequentials() (in module penzai.deprecated.v1.nn.grouping) (in module penzai.nn.grouping) inline_groups() (in module penzai.deprecated.v1.nn.grouping) (in module penzai.nn.grouping) input_axes (penzai.deprecated.v1.nn.linear_and_affine.Linear property) (penzai.nn.linear_and_affine.Linear property) input_structure() (penzai.deprecated.v1.core.layer.Layer method) insert_after() (penzai.core.selectors.Selection method) insert_before() (penzai.core.selectors.Selection method) InterceptedFlaxModuleMethod (class in penzai.deprecated.v1.toolshed.unflaxify) (class in penzai.toolshed.unflaxify) InterceptedFlaxScopeData (class in penzai.deprecated.v1.toolshed.unflaxify) (class in penzai.toolshed.unflaxify) InternalTrainerState (class in penzai.toolshed.basic_training) into_pytree() (penzai.core.shapecheck.ArraySpec method) invert() (penzai.core.selectors.Selection method) is_anonymous_sequential() (in module penzai.deprecated.v1.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.deprecated.v1.nn.grouping) (in module penzai.nn.grouping) IsControl() (penzai.toolshed.token_visualization.SentencePieceProcessorLike method) IsolatedSubmodel (class in penzai.deprecated.v1.toolshed.isolate_submodel) (class in penzai.toolshed.isolate_submodel) IsolationCapturer (class in penzai.deprecated.v1.toolshed.isolate_submodel) (class in penzai.toolshed.isolate_submodel) item() (penzai.core.named_axes.NamedArrayBase method) J Jitted (class in penzai.deprecated.v1.toolshed.jit_wrapper) (class in penzai.toolshed.jit_wrapper) K key_for_field() (penzai.core.struct.Struct method) (penzai.nn.layer_stack.LayerStack method) KnockOutAttentionHeads (class in penzai.deprecated.v1.toolshed.model_rewiring) (class in penzai.toolshed.model_rewiring) KnownDim (class in penzai.core.shapecheck) KVCachingAttention (class in penzai.deprecated.v1.nn.attention) (class in penzai.nn.attention) KVCachingTransformerLM (class in penzai.models.transformer.sampling_mode) L LabeledVariable (class in penzai.core.variables) LabeledVariableValue (class in penzai.core.variables) Layer (class in penzai.deprecated.v1.core.layer) (class in penzai.nn.layer) LayerNorm (class in penzai.deprecated.v1.nn.standardization) (class in penzai.nn.standardization) LayerStack (class in penzai.nn.layer_stack) layerstack_axes_from_keypath() (in module penzai.nn.layer_stack) LayerStackGetAttrKey (class in penzai.nn.layer_stack) LayerStackVarBehavior (class in penzai.nn.layer_stack) Linear (class in penzai.deprecated.v1.nn.linear_and_affine) (class in penzai.nn.linear_and_affine) LinearInPlace (class in penzai.deprecated.v1.nn.linear_and_affine) (class in penzai.nn.linear_and_affine) LinearizeAndAdjust (class in penzai.deprecated.v1.toolshed.model_rewiring) (class in penzai.toolshed.model_rewiring) LinearOperatorWeightInitializer (class in penzai.deprecated.v1.nn.linear_and_affine) (class in penzai.nn.linear_and_affine) llama_from_huggingface_model() (in module penzai.models.transformer.variants.llama) llamalike_from_huggingface_model() (in module penzai.models.transformer.variants.llamalike_common) LlamalikeTransformerConfig (class in penzai.models.transformer.variants.llamalike_common) LocalStateEffect (class in penzai.deprecated.v1.data_effects.local_state) LocalStateEffectImpl (class in penzai.deprecated.v1.data_effects.local_state) loraify_linears_in_selection() (in module penzai.deprecated.v1.toolshed.lora) (in module penzai.toolshed.lora) LossFunction (class in penzai.deprecated.v1.toolshed.basic_training) (class in penzai.toolshed.basic_training) LowRankAdapter (class in penzai.deprecated.v1.toolshed.lora) (class in penzai.toolshed.lora) M make_parameter() (in module penzai.nn.parameters) mark_shareable() (in module penzai.deprecated.v1.nn.parameters) MatchResult (class in penzai.core.shapecheck) max() (penzai.core.named_axes.NamedArrayBase method) maybe_dropout() (in module penzai.deprecated.v1.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.models.transformer.variants.mistral) MLP (class in penzai.deprecated.v1.example_models.simple_mlp) (class in penzai.models.simple_mlp) model (penzai.deprecated.v1.toolshed.basic_training.TrainState property) module penzai.core penzai.core.auto_order_types 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.core.variables penzai.deprecated.v1 penzai.deprecated.v1.core penzai.deprecated.v1.core.layer penzai.deprecated.v1.core.random_stream penzai.deprecated.v1.data_effects penzai.deprecated.v1.data_effects.effect_base penzai.deprecated.v1.data_effects.local_state penzai.deprecated.v1.data_effects.random penzai.deprecated.v1.data_effects.side_input penzai.deprecated.v1.data_effects.side_output penzai.deprecated.v1.example_models penzai.deprecated.v1.example_models.gemma penzai.deprecated.v1.example_models.gemma.model_core penzai.deprecated.v1.example_models.gemma.sampling_mode penzai.deprecated.v1.example_models.gemma.simple_decoding_loop penzai.deprecated.v1.example_models.simple_mlp penzai.deprecated.v1.nn penzai.deprecated.v1.nn.attention penzai.deprecated.v1.nn.basic_ops penzai.deprecated.v1.nn.combinators penzai.deprecated.v1.nn.dropout penzai.deprecated.v1.nn.embeddings penzai.deprecated.v1.nn.grouping penzai.deprecated.v1.nn.linear_and_affine penzai.deprecated.v1.nn.parameters penzai.deprecated.v1.nn.standardization penzai.deprecated.v1.toolshed penzai.deprecated.v1.toolshed.annotate_shapes penzai.deprecated.v1.toolshed.basic_training penzai.deprecated.v1.toolshed.check_layers_by_tracing penzai.deprecated.v1.toolshed.interleave_intermediates penzai.deprecated.v1.toolshed.isolate_submodel penzai.deprecated.v1.toolshed.jit_wrapper penzai.deprecated.v1.toolshed.lora penzai.deprecated.v1.toolshed.model_rewiring penzai.deprecated.v1.toolshed.sharding_util penzai.deprecated.v1.toolshed.unflaxify penzai.models penzai.models.simple_mlp penzai.models.transformer penzai.models.transformer.model_parts penzai.models.transformer.sampling_mode penzai.models.transformer.simple_decoding_loop penzai.models.transformer.variants penzai.models.transformer.variants.gemma penzai.models.transformer.variants.gpt_neox penzai.models.transformer.variants.llama penzai.models.transformer.variants.llamalike_common penzai.models.transformer.variants.mistral penzai.nn penzai.nn.attention penzai.nn.basic_ops penzai.nn.combinators penzai.nn.dropout penzai.nn.embeddings penzai.nn.grouping penzai.nn.layer penzai.nn.layer_stack penzai.nn.linear_and_affine penzai.nn.parameters penzai.nn.standardization penzai.pz penzai.pz.nn penzai.pz.ts penzai.toolshed penzai.toolshed.auto_nmap penzai.toolshed.basic_training penzai.toolshed.gradient_checkpointing penzai.toolshed.isolate_submodel penzai.toolshed.jit_wrapper penzai.toolshed.lora penzai.toolshed.model_rewiring penzai.toolshed.patch_ipdb penzai.toolshed.save_intermediates penzai.toolshed.sharding_util penzai.toolshed.token_visualization penzai.toolshed.unflaxify penzai.treescope 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.deprecated.v1.nn.linear_and_affine) (class in penzai.nn.linear_and_affine) NamedGroup (class in penzai.deprecated.v1.nn.grouping) (class in penzai.nn.grouping) next_key() (penzai.core.random_stream.RandomStream method) (penzai.deprecated.v1.core.random_stream.RandomStream method) (penzai.deprecated.v1.data_effects.random.RandomEffect method) nmap() (in module penzai.core.named_axes) nonzero() (penzai.core.named_axes.NamedArrayBase method) NotInThisPartition (class in penzai.core.partitioning) O 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.deprecated.v1.nn.linear_and_affine.Linear property) (penzai.nn.linear_and_affine.Linear property) output_structure() (penzai.deprecated.v1.core.layer.Layer method) P parallel_axes (penzai.deprecated.v1.nn.linear_and_affine.Linear property) (penzai.nn.linear_and_affine.Linear property) Parameter (class in penzai.core.variables) (class in penzai.deprecated.v1.nn.parameters) ParameterLike (class in penzai.deprecated.v1.nn.parameters) (class in penzai.nn.parameters) ParameterSlot (class in penzai.core.variables) ParameterValue (class in penzai.core.variables) partition() (penzai.core.selectors.Selection method) patch_ipdb() (in module penzai.toolshed.patch_ipdb) path_matrix() (penzai.deprecated.v1.toolshed.model_rewiring.RewireComputationPaths method) (penzai.toolshed.model_rewiring.RewireComputationPaths method) penzai.core module penzai.core.auto_order_types 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.core.variables module penzai.deprecated.v1 module penzai.deprecated.v1.core module penzai.deprecated.v1.core.layer module penzai.deprecated.v1.core.random_stream module penzai.deprecated.v1.data_effects module penzai.deprecated.v1.data_effects.effect_base module penzai.deprecated.v1.data_effects.local_state module penzai.deprecated.v1.data_effects.random module penzai.deprecated.v1.data_effects.side_input module penzai.deprecated.v1.data_effects.side_output module penzai.deprecated.v1.example_models module penzai.deprecated.v1.example_models.gemma module penzai.deprecated.v1.example_models.gemma.model_core module penzai.deprecated.v1.example_models.gemma.sampling_mode module penzai.deprecated.v1.example_models.gemma.simple_decoding_loop module penzai.deprecated.v1.example_models.simple_mlp module penzai.deprecated.v1.nn module penzai.deprecated.v1.nn.attention module penzai.deprecated.v1.nn.basic_ops module penzai.deprecated.v1.nn.combinators module penzai.deprecated.v1.nn.dropout module penzai.deprecated.v1.nn.embeddings module penzai.deprecated.v1.nn.grouping module penzai.deprecated.v1.nn.linear_and_affine module penzai.deprecated.v1.nn.parameters module penzai.deprecated.v1.nn.standardization module penzai.deprecated.v1.toolshed module penzai.deprecated.v1.toolshed.annotate_shapes module penzai.deprecated.v1.toolshed.basic_training module penzai.deprecated.v1.toolshed.check_layers_by_tracing module penzai.deprecated.v1.toolshed.interleave_intermediates module penzai.deprecated.v1.toolshed.isolate_submodel module penzai.deprecated.v1.toolshed.jit_wrapper module penzai.deprecated.v1.toolshed.lora module penzai.deprecated.v1.toolshed.model_rewiring module penzai.deprecated.v1.toolshed.sharding_util module penzai.deprecated.v1.toolshed.unflaxify module penzai.models module penzai.models.simple_mlp module penzai.models.transformer module penzai.models.transformer.model_parts module penzai.models.transformer.sampling_mode module penzai.models.transformer.simple_decoding_loop module penzai.models.transformer.variants module penzai.models.transformer.variants.gemma module penzai.models.transformer.variants.gpt_neox module penzai.models.transformer.variants.llama module penzai.models.transformer.variants.llamalike_common module penzai.models.transformer.variants.mistral 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.layer module penzai.nn.layer_stack module penzai.nn.linear_and_affine module penzai.nn.parameters module penzai.nn.standardization module penzai.pz module penzai.pz.nn module penzai.pz.ts module penzai.toolshed module penzai.toolshed.auto_nmap module penzai.toolshed.basic_training module penzai.toolshed.gradient_checkpointing 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.save_intermediates module penzai.toolshed.sharding_util module penzai.toolshed.token_visualization module penzai.toolshed.unflaxify module penzai.treescope module pick_nth_selected() (penzai.core.selectors.Selection method) positional_shape (penzai.core.named_axes.NamedArrayBase property) prefill() (in module penzai.deprecated.v1.example_models.gemma.simple_decoding_loop) pretty_keystr() (in module penzai.core.tree_util) 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.deprecated.v1.data_effects.random) RandomEffectImpl (class in penzai.deprecated.v1.data_effects.random) RandomRequest (class in penzai.deprecated.v1.data_effects.random) RandomStream (class in penzai.core.random_stream) (class in penzai.deprecated.v1.core.random_stream) ravel() (penzai.core.named_axes.NamedArrayBase method) real (penzai.core.named_axes.NamedArrayBase attribute) refine() (penzai.core.selectors.Selection method) register_effect_color() (in module penzai.deprecated.v1.data_effects.effect_base) RemainingAxisPlaceholder (class in penzai.core.shapecheck) remove_from_parent() (penzai.core.selectors.Selection method) RenameAxes (class in penzai.deprecated.v1.nn.linear_and_affine) (class in penzai.nn.linear_and_affine) repeat() (penzai.core.named_axes.NamedArrayBase method) reshape() (penzai.core.named_axes.NamedArrayBase method) Residual (class in penzai.deprecated.v1.nn.combinators) (class in penzai.nn.combinators) RewireComputationPaths (class in penzai.deprecated.v1.toolshed.model_rewiring) (class in penzai.toolshed.model_rewiring) RMSLayerNorm (class in penzai.deprecated.v1.nn.standardization) (class in penzai.nn.standardization) RMSStandardize (class in penzai.deprecated.v1.nn.standardization) (class in penzai.nn.standardization) round() (penzai.core.named_axes.NamedArrayBase method) run_and_interleave_intermediates() (in module penzai.deprecated.v1.toolshed.interleave_intermediates) S SamplingState (class in penzai.deprecated.v1.example_models.gemma.simple_decoding_loop) SaveIntermediate (class in penzai.toolshed.save_intermediates) SaveIntermediateShape (class in penzai.toolshed.save_intermediates) saving_all_intermediates() (in module penzai.toolshed.save_intermediates) scan() (in module penzai.core.named_axes) scoped_auto_state_var_labels() (in module penzai.core.variables) ScopedStateVarLabel (class in penzai.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.deprecated.v1.nn.grouping) (class in penzai.nn.grouping) set() (penzai.core.selectors.Selection method) (penzai.deprecated.v1.data_effects.local_state.LocalStateEffect method) set_by_path() (penzai.core.selectors.Selection method) set_sequence() (penzai.core.selectors.Selection method) set_value() (penzai.core.variables.LabeledVariable method) sharded_init() (in module penzai.toolshed.sharding_util) ShareableUninitializedParameter (class in penzai.deprecated.v1.nn.parameters) SharedLocalStateRequest (class in penzai.deprecated.v1.data_effects.local_state) SharedParameterLookup (class in penzai.deprecated.v1.nn.parameters) SharedParamTag (class in penzai.deprecated.v1.nn.parameters) 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.deprecated.v1.data_effects.side_input) SideInputEffectImpl (class in penzai.deprecated.v1.data_effects.side_input) SideInputRequest (class in penzai.deprecated.v1.data_effects.side_input) SideOutputEffect (class in penzai.deprecated.v1.data_effects.side_output) SideOutputEffectImpl (class in penzai.deprecated.v1.data_effects.side_output) SideOutputRequest (class in penzai.deprecated.v1.data_effects.side_output) SideOutputValue (class in penzai.deprecated.v1.data_effects.side_output) SliceLike (class in penzai.core.syntactic_sugar) Softmax (class in penzai.deprecated.v1.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.deprecated.v1.nn.standardization) (class in penzai.nn.standardization) StatefulTrainer (class in penzai.toolshed.basic_training) stateless_call() (penzai.nn.layer.Layer method) StateVariable (class in penzai.core.variables) StateVariableSlot (class in penzai.core.variables) StateVariableValue (class in penzai.core.variables) Static (class in penzai.deprecated.v1.toolshed.annotate_shapes) std() (penzai.core.named_axes.NamedArrayBase method) step() (penzai.toolshed.basic_training.StatefulTrainer method) StochasticDropout (class in penzai.deprecated.v1.nn.dropout) (class in penzai.nn.dropout) Struct (class in penzai.core.struct) StructStaticMetadata (class in penzai.core.struct) StructureMismatchError sum() (penzai.core.named_axes.NamedArrayBase method) SupportsParameterRenaming (class in penzai.deprecated.v1.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.deprecated.v1.data_effects.random) take() (penzai.core.named_axes.NamedArrayBase method) TanhSoftCap (class in penzai.nn.basic_ops) tell() (penzai.deprecated.v1.data_effects.side_output.SideOutputEffect method) TellIntermediate (class in penzai.deprecated.v1.data_effects.side_output) temperature_sample_pyloop() (in module penzai.deprecated.v1.example_models.gemma.simple_decoding_loop) (in module penzai.models.transformer.simple_decoding_loop) TmpPosAxisMarker (class in penzai.core.named_axes) trace() (penzai.core.named_axes.NamedArrayBase method) TrainState (class in penzai.deprecated.v1.toolshed.basic_training) TrainStepFunction (class in penzai.deprecated.v1.toolshed.basic_training) TransformerBlock (class in penzai.models.transformer.model_parts) TransformerFeedForward (class in penzai.models.transformer.model_parts) TransformerLM (class in penzai.models.transformer.model_parts) TransformerMetadata (class in penzai.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.core.variables.LabeledVariable method) U unbind_params() (in module penzai.core.variables) unbind_state_vars() (in module penzai.core.variables) unbind_variables() (in module penzai.core.variables) UnboundVariableError unchecked_layer_call() (in module penzai.deprecated.v1.core.layer) unflaxify_apply() (in module penzai.deprecated.v1.toolshed.unflaxify) (in module penzai.toolshed.unflaxify) unfreeze_as_copy() (penzai.core.variables.AbstractVariableValue method) UnhandledEffectError UninitializedParameter (class in penzai.deprecated.v1.nn.parameters) UninitializedParameterError unpatch_ipdb() (in module penzai.toolshed.patch_ipdb) unsafe_mark_active() (penzai.deprecated.v1.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.core.variables.AbstractVariable method) (penzai.core.variables.LabeledVariable method) V value (penzai.core.variables.ParameterSlot property) (penzai.core.variables.StateVariableSlot property) (penzai.deprecated.v1.nn.parameters.ParameterLike property) (penzai.deprecated.v1.nn.parameters.SharedParameterLookup property) (penzai.deprecated.v1.nn.parameters.UninitializedParameter property) (penzai.nn.parameters.ParameterLike property) value_structure (penzai.deprecated.v1.nn.parameters.ParameterLike property) var() (in module penzai.core.shapecheck) (penzai.core.named_axes.NamedArrayBase method) variable_jit() (in module penzai.core.variables) VariableConflictError variance_scaling_initializer() (in module penzai.deprecated.v1.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_positional_prefix() (penzai.core.named_axes.NamedArrayBase method) (penzai.core.named_axes.NamedArrayView method) with_renamed_parameters() (penzai.deprecated.v1.nn.parameters.SupportsParameterRenaming method) WithConstantSideInputs (class in penzai.deprecated.v1.data_effects.side_input) WithFrozenRandomState (class in penzai.deprecated.v1.data_effects.random) WithFunctionalLocalState (class in penzai.deprecated.v1.data_effects.local_state) WithRandomKeyFromArg (class in penzai.deprecated.v1.data_effects.random) WithSideInputsFromInputTuple (class in penzai.deprecated.v1.data_effects.side_input) WithStatefulRandomKey (class in penzai.deprecated.v1.data_effects.random) wrap() (in module penzai.core.named_axes) (penzai.core.named_axes.NamedArray class method) wrap_module() (in module penzai.toolshed.auto_nmap) Z zero_initializer() (in module penzai.deprecated.v1.nn.linear_and_affine) (in module penzai.nn.linear_and_affine) zeros() (in module penzai.core.named_axes)