penzai.deprecated.v1 (Original V1 API)

penzai.deprecated.v1 is Penzai’s original neural network system, which has been replaced by the new “V2” design. It is available here for compatibility with existing users of the V1 API. For a summary of the differences between the V1 and V2 APIs, and a guide to migrating to the V2 API, see the “Changes in the V2 API” overview document.

To use the V1 API, we suggest importing the pz alias namespace from penzai.deprecated.v1.pz:

from penzai.deprecated.v1 import pz

The rest of this page lists the main components used in the V1 API.

Specific to the V1 Neural Network API

V1 Layers

Most models and layers in the V1 API are subclasses of pz.Layer, which means they can be called with a single argument.

v1.pz.Layer	Alias of penzai.deprecated.v1.core.layer.Layer: Abstract base class for neural network layers and other 1-arg callables.
v1.pz.LayerLike	alias of Callable[[Any], Any]

The following decorators can be used to enable runtime shape-checking on penzai.deprecated.v1.core.layer.Layer (pz.Layer) subclasses:

v1.pz.checked_layer_call(func)	Alias of penzai.deprecated.v1.core.layer.checked_layer_call: Decorator for Layer.__call__ to add shape-checking and name scopes.
v1.pz.unchecked_layer_call(func)	Alias of penzai.deprecated.v1.core.layer.unchecked_layer_call: Shape-checking opt-out decorator for Layer.__call__.

V1 Parameters

In the V1 API, Parameters are always ordinary PyTree nodes. Parameter sharing is indicated by extra metadata and PyTree transformations inside the model.

v1.pz.nn.Parameter	Alias of penzai.deprecated.v1.nn.parameters.Parameter: A learnable parameter.
v1.pz.nn.ParameterLike	Alias of penzai.deprecated.v1.nn.parameters.ParameterLike: Protocol for a parameter-like object.
v1.pz.nn.UninitializedParameter	Alias of penzai.deprecated.v1.nn.parameters.UninitializedParameter: An uninitialized parameter.
v1.pz.nn.add_parameter_prefix(prefix, tree)	Alias of penzai.deprecated.v1.nn.parameters.add_parameter_prefix: Prepends a prefix to all parameter names inside the tree.
v1.pz.nn.initialize_parameters(model, prng_key)	Alias of penzai.deprecated.v1.nn.parameters.initialize_parameters: Initializes all parameters in a model.
v1.pz.nn.FrozenParameter	Alias of penzai.deprecated.v1.nn.parameters.FrozenParameter: A non-learnable parameter.
v1.pz.nn.mark_shareable(submodel)	Alias of penzai.deprecated.v1.nn.parameters.mark_shareable: Marks all uninitialized parameters in submodel as shareable.
v1.pz.nn.ShareableUninitializedParameter	Alias of penzai.deprecated.v1.nn.parameters.ShareableUninitializedParameter: A shareable variant of an uninitialized parameter.
v1.pz.nn.attach_shared_parameters(submodel, *)	Alias of penzai.deprecated.v1.nn.parameters.attach_shared_parameters: Wraps a submodel in a handler that takes ownership of shareable params.
v1.pz.nn.SharedParameterLookup	Alias of penzai.deprecated.v1.nn.parameters.SharedParameterLookup: A marker identifying a shared parameter.
v1.pz.nn.SharedParamTag	Alias of penzai.deprecated.v1.nn.parameters.SharedParamTag: A tag for a shared parameter.
v1.pz.nn.SupportsParameterRenaming	Alias of penzai.deprecated.v1.nn.parameters.SupportsParameterRenaming: Base class that identifies a PyTree node as supporting parameter renaming.
v1.pz.nn.check_no_duplicated_parameters(model)	Alias of penzai.deprecated.v1.nn.parameters.check_no_duplicated_parameters: Checks that there are no duplicated parameters in a model.
v1.pz.nn.UninitializedParameterError	Alias of penzai.deprecated.v1.nn.parameters.UninitializedParameterError: Raised when an uninitialized parameter is accessed.

Basic Combinators

v1.pz.nn.Sequential	Alias of penzai.deprecated.v1.nn.grouping.Sequential: A group of layers to call sequentially.
v1.pz.nn.NamedGroup	Alias of penzai.deprecated.v1.nn.grouping.NamedGroup: A layer that names an activation or a sequence of layers.
v1.pz.nn.CheckedSequential	Alias of penzai.deprecated.v1.nn.grouping.CheckedSequential: A group of layers to call sequentially, with known input/output types.
v1.pz.nn.Residual	Alias of penzai.deprecated.v1.nn.combinators.Residual: A residual block, which runs its sublayers then adds the input.
v1.pz.nn.BranchAndAddTogether	Alias of penzai.deprecated.v1.nn.combinators.BranchAndAddTogether: A data-flow branch with additive interactions between branches.
v1.pz.nn.BranchAndMultiplyTogether	Alias of penzai.deprecated.v1.nn.combinators.BranchAndMultiplyTogether: A data-flow branch with multiplicative interactions between branches.
v1.pz.nn.inline_anonymous_sequentials(tree)	Alias of penzai.deprecated.v1.nn.grouping.inline_anonymous_sequentials: Inlines instances of Sequential (not subclasses) into parent groups.
v1.pz.nn.inline_groups(tree, parent_filter, ...)	Alias of penzai.deprecated.v1.nn.grouping.inline_groups: Inlines sequential nodes into their parents if possible.
v1.pz.nn.is_anonymous_sequential(tree)	Alias of penzai.deprecated.v1.nn.grouping.is_anonymous_sequential: Checks if the type of a node is exactly Sequential, not a named subclass.
v1.pz.nn.is_sequential_or_named(tree)	Alias of penzai.deprecated.v1.nn.grouping.is_sequential_or_named: Checks if a tree is a subclass of Sequential or a NamedGroup.

Basic Operations

v1.pz.nn.Elementwise	Alias of penzai.deprecated.v1.nn.basic_ops.Elementwise: A layer that runs an elementwise operation on its NamedArray argument.
v1.pz.nn.Softmax	Alias of penzai.deprecated.v1.nn.basic_ops.Softmax: Layer that applies a softmax along a given set of axes.
v1.pz.nn.CheckStructure	Alias of penzai.deprecated.v1.nn.grouping.CheckStructure: A layer that checks the structure of the value passing through it.
v1.pz.nn.Identity	Alias of penzai.deprecated.v1.nn.grouping.Identity: A layer that returns its input unchanged, without any side effects.
v1.pz.nn.CastToDType	Alias of penzai.deprecated.v1.nn.basic_ops.CastToDType: Casts an input to a given dtype.

Linear and Affine Layers

v1.pz.nn.Linear	Alias of penzai.deprecated.v1.nn.linear_and_affine.Linear: A generalized linear (not affine) operator, for named arrays.
v1.pz.nn.RenameAxes	Alias of penzai.deprecated.v1.nn.linear_and_affine.RenameAxes: Convenience layer that renames axes of its input.
v1.pz.nn.AddBias	Alias of penzai.deprecated.v1.nn.linear_and_affine.AddBias: Shifts an input by a learnable offset (a bias).
v1.pz.nn.Affine	Alias of penzai.deprecated.v1.nn.linear_and_affine.Affine: Affine layer: combination of Linear and AddBias.
v1.pz.nn.ConstantRescale	Alias of penzai.deprecated.v1.nn.linear_and_affine.ConstantRescale: Applies a constant scaling factor.
v1.pz.nn.NamedEinsum	Alias of penzai.deprecated.v1.nn.linear_and_affine.NamedEinsum: An Einsum operation that contracts based on axis names.
v1.pz.nn.LinearInPlace	Alias of penzai.deprecated.v1.nn.linear_and_affine.LinearInPlace: Container for "in-place" linear operators that preserve axis names.
v1.pz.nn.LinearOperatorWeightInitializer	Alias of penzai.deprecated.v1.nn.linear_and_affine.LinearOperatorWeightInitializer: Protocol for an initializer for a general linear NamedArray weight.
v1.pz.nn.BiasInitializer	alias of Callable[[Array, tuple[int, ...], dtype], Array]
v1.pz.nn.contract(names, left, right)	Alias of penzai.deprecated.v1.nn.linear_and_affine.contract: Contracts two named arrays along the given axis names.
v1.pz.nn.variance_scaling_initializer(key, ...)	Alias of penzai.deprecated.v1.nn.linear_and_affine.variance_scaling_initializer: Generic variance scaling initializer.
v1.pz.nn.xavier_normal_initializer(key, *[, ...])	Generic variance scaling initializer.
v1.pz.nn.xavier_uniform_initializer(key, *)	Generic variance scaling initializer.
v1.pz.nn.constant_initializer(value)	Alias of penzai.deprecated.v1.nn.linear_and_affine.constant_initializer: Returns an initializer that uses a constant value.
v1.pz.nn.zero_initializer(key, *, ...)	Zeros initializer for named arrays.

Standardization

v1.pz.nn.LayerNorm	Alias of penzai.deprecated.v1.nn.standardization.LayerNorm: Layer normalization layer.
v1.pz.nn.Standardize	Alias of penzai.deprecated.v1.nn.standardization.Standardize: Standardization layer.
v1.pz.nn.RMSLayerNorm	Alias of penzai.deprecated.v1.nn.standardization.RMSLayerNorm: Root-mean-squared layer normalization layer.
v1.pz.nn.RMSStandardize	Alias of penzai.deprecated.v1.nn.standardization.RMSStandardize: Root-mean-squared standardization layer.

Dropout

v1.pz.nn.StochasticDropout	Alias of penzai.deprecated.v1.nn.dropout.StochasticDropout: Stochastic dropout layer.
v1.pz.nn.DisabledDropout	Alias of penzai.deprecated.v1.nn.dropout.DisabledDropout: A no-op layer taking the place of a disabled StochasticDropout layer.
v1.pz.nn.maybe_dropout(drop_rate[, ...])	Alias of penzai.deprecated.v1.nn.dropout.maybe_dropout: Constructs either a stochastic or disabled dropout layer.

Language Modeling

v1.pz.nn.Attention	Alias of penzai.deprecated.v1.nn.attention.Attention: A basic attention combinator.
v1.pz.nn.KVCachingAttention	Alias of penzai.deprecated.v1.nn.attention.KVCachingAttention: Key/value caching variant of Attention.
v1.pz.nn.ApplyAttentionMask	Alias of penzai.deprecated.v1.nn.attention.ApplyAttentionMask: Applies an attention mask to its input logit array.
v1.pz.nn.EmbeddingTable	Alias of penzai.deprecated.v1.nn.embeddings.EmbeddingTable: A table of embedding vectors for a vocabulary of tokens.
v1.pz.nn.EmbeddingLookup	Alias of penzai.deprecated.v1.nn.embeddings.EmbeddingLookup: Looks up token IDs in an embedding table.
v1.pz.nn.EmbeddingDecode	Alias of penzai.deprecated.v1.nn.embeddings.EmbeddingDecode: Uses an embedding table to map embeddings back to token scores.
v1.pz.nn.ApplyRoPE	Alias of penzai.deprecated.v1.nn.embeddings.ApplyRoPE: Adjusts input embeddings using rotary position embeddings (RoPE).

Base Effect Types

v1.pz.de.EffectHandler	Alias of penzai.deprecated.v1.data_effects.effect_base.EffectHandler: A handler for a particular effect.
v1.pz.de.EffectRequest	Alias of penzai.deprecated.v1.data_effects.effect_base.EffectRequest: Base class for "effect requests", which represent unhandled effects.
v1.pz.de.EffectRuntimeImpl	Alias of penzai.deprecated.v1.data_effects.effect_base.EffectRuntimeImpl: Base class for runtime effect implementations.
v1.pz.de.HandledEffectRef	Alias of penzai.deprecated.v1.data_effects.effect_base.HandledEffectRef: Base class for references to a handler that handles this effect.
v1.pz.de.HandlerId	alias of str

Side Inputs

v1.pz.de.SideInputEffect	Alias of penzai.deprecated.v1.data_effects.side_input.SideInputEffect: Protocol for a side input effect.
v1.pz.de.SideInputRequest	Alias of penzai.deprecated.v1.data_effects.side_input.SideInputRequest: Effect request for a side input.
v1.pz.de.HandledSideInputRef	Alias of penzai.deprecated.v1.data_effects.side_input.HandledSideInputRef: Reference for a handled side input effect.
v1.pz.de.WithSideInputsFromInputTuple	Alias of penzai.deprecated.v1.data_effects.side_input.WithSideInputsFromInputTuple: SideInput handler that unpacks side inputs from a tuple argument.
v1.pz.de.WithConstantSideInputs	Alias of penzai.deprecated.v1.data_effects.side_input.WithConstantSideInputs: SideInput handler that provides side inputs using its own attribute.
v1.pz.de.HoistedTag	Alias of penzai.deprecated.v1.data_effects.side_input.HoistedTag: A tag that has been hoisted out of a handler.
v1.pz.de.hoist_constant_side_inputs(tree[, ...])	Alias of penzai.deprecated.v1.data_effects.side_input.hoist_constant_side_inputs: Extracts all constant side inputs from a tree so they can be re-handled.

Side Outputs

v1.pz.de.CollectingSideOutputs	Alias of penzai.deprecated.v1.data_effects.side_output.CollectingSideOutputs: SideOutput handler that collects all side outputs into a list.
v1.pz.de.HandledSideOutputRef	Alias of penzai.deprecated.v1.data_effects.side_output.HandledSideOutputRef: Marker for a handled local state effect.
v1.pz.de.SideOutputEffect	Alias of penzai.deprecated.v1.data_effects.side_output.SideOutputEffect: Protocol for a side output effect.
v1.pz.de.SideOutputRequest	Alias of penzai.deprecated.v1.data_effects.side_output.SideOutputRequest: Effect request for a side output.
v1.pz.de.SideOutputValue	Alias of penzai.deprecated.v1.data_effects.side_output.SideOutputValue: A value written to a side output.
v1.pz.de.TellIntermediate	Alias of penzai.deprecated.v1.data_effects.side_output.TellIntermediate: Helper layer that writes its intermediate value to a side output.

Randomness

v1.pz.de.RandomEffect	Alias of penzai.deprecated.v1.data_effects.random.RandomEffect: Protocol for the random number generation effect.
v1.pz.de.RandomRequest	Alias of penzai.deprecated.v1.data_effects.random.RandomRequest: Random number generation request.
v1.pz.de.TaggedRandomRequest	Alias of penzai.deprecated.v1.data_effects.random.TaggedRandomRequest: Random number generation request, with a tag.
v1.pz.de.HandledRandomRef	Alias of penzai.deprecated.v1.data_effects.random.HandledRandomRef: Reference for a handled random effect.
v1.pz.de.WithRandomKeyFromArg	Alias of penzai.deprecated.v1.data_effects.random.WithRandomKeyFromArg: RandomEffect handler that expects a random seed as its second argument.
v1.pz.de.WithStatefulRandomKey	Alias of penzai.deprecated.v1.data_effects.random.WithStatefulRandomKey: RandomEffect handler that tracks a random seed as a local state.
v1.pz.de.WithFrozenRandomState	Alias of penzai.deprecated.v1.data_effects.random.WithFrozenRandomState: RandomEffect handler that uses a fixed random state.

Local State

v1.pz.de.LocalStateEffect	Alias of penzai.deprecated.v1.data_effects.local_state.LocalStateEffect: Protocol for a local state effect.
v1.pz.de.InitialLocalStateRequest	Alias of penzai.deprecated.v1.data_effects.local_state.InitialLocalStateRequest: Effect request for local state, with a state initializer.
v1.pz.de.FrozenLocalStateRequest	Alias of penzai.deprecated.v1.data_effects.local_state.FrozenLocalStateRequest: Effect request for local state with a frozen value.
v1.pz.de.SharedLocalStateRequest	Alias of penzai.deprecated.v1.data_effects.local_state.SharedLocalStateRequest: Effect request for local state that is shared.
v1.pz.de.HandledLocalStateRef	Alias of penzai.deprecated.v1.data_effects.local_state.HandledLocalStateRef: Marker for a handled local state effect.
v1.pz.de.WithFunctionalLocalState	Alias of penzai.deprecated.v1.data_effects.local_state.WithFunctionalLocalState: LocalState effect handler that functionalizes local states.
v1.pz.de.handle_local_states()	Alias of penzai.deprecated.v1.data_effects.local_state.handle_local_states: Extracts local states from a stateful model.
v1.pz.de.freeze_local_states(handled, states)	Alias of penzai.deprecated.v1.data_effects.local_state.freeze_local_states: Embeds the given states into a handled model, and removes the handler.
v1.pz.de.hoist_shared_state_requests(tree[, ...])	Alias of penzai.deprecated.v1.data_effects.local_state.hoist_shared_state_requests: Hoists out the value of shared states in a pytree.
v1.pz.de.embed_shared_state_requests(tree, ...)	Alias of penzai.deprecated.v1.data_effects.local_state.embed_shared_state_requests: Embeds shared state requests into a tree.

Effect Utilities

v1.pz.de.all_handler_ids(model_tree)	Alias of penzai.deprecated.v1.data_effects.effect_base.all_handler_ids: Collects the set of all handler IDs inside a model or submodel.
v1.pz.de.free_effect_types(model_tree)	Alias of penzai.deprecated.v1.data_effects.effect_base.free_effect_types: Collects the effect types of all EffectRequest nodes in a (sub)model.
v1.pz.de.get_effect_color(effect_protocol)	Alias of penzai.deprecated.v1.data_effects.effect_base.get_effect_color: Gets the default color for a given effect (for treescope rendering).
v1.pz.de.infer_or_check_handler_id(tag, subtree)	Alias of penzai.deprecated.v1.data_effects.effect_base.infer_or_check_handler_id: Tries to generate a unique handler ID from the structure of a subtree.
v1.pz.de.register_effect_color(color)	Alias of penzai.deprecated.v1.data_effects.effect_base.register_effect_color: Decorator to register a treescope-rendering color for a given effect.
v1.pz.de.UnhandledEffectError	Alias of penzai.deprecated.v1.data_effects.effect_base.UnhandledEffectError: Exception raised when a method is called on an unhandled effect.

Core utilities, shared with the V2 API

Structs and Layers

Most objects in Penzai models are subclasses of pz.Struct and decorated with pz.pytree_dataclass, which makes them into frozen Python dataclasses that are also JAX PyTrees.

v1.pz.pytree_dataclass([cls, ...])	Alias of penzai.core.struct.pytree_dataclass: Decorator for constructing a frozen PyTree dataclass.
v1.pz.Struct	Alias of penzai.core.struct.Struct: Base class for penzai PyTree structures.

PyTree Manipulation

Penzai provides a number of utilities to make targeted modifications to PyTrees. Since Penzai models are PyTrees, you can use them to insert new layers into models, or modify the configuration of existing layers.

v1.pz.select(tree)	Alias of penzai.core.selectors.select: Wraps a tree in a singleton selection for processing.
v1.pz.Selection	Alias of penzai.core.selectors.Selection: A selected subset of nodes within a larger PyTree.
v1.pz.combine(*partitions)	Alias of penzai.core.partitioning.combine: Combines leaves from multiple partitions.
v1.pz.NotInThisPartition	Alias of penzai.core.partitioning.NotInThisPartition: Sentinel object that identifies subtrees removed by partition.
v1.pz.pretty_keystr(keypath, tree)	Alias of penzai.core.tree_util.pretty_keystr: Constructs a pretty name from a keypath and an object.

Named Axes

pz.nx is an alias for penzai.core.named_axes, which contains Penzai’s named axis system. Some commonly-used attributes on pz.nx:

v1.pz.nx.NamedArray	Alias of penzai.core.named_axes.NamedArray: A multidimensional array with a combination of positional and named axes.
v1.pz.nx.nmap(fun)	Alias of penzai.core.named_axes.nmap: Automatically vectorizes fun over named axes of NamedArray inputs.
v1.pz.nx.wrap(array, *names)	Alias of penzai.core.named_axes.NamedArray.wrap: Wraps a positional array as a NamedArray.

See penzai.core.named_axes for documentation of all of the methods and classes accessible through the pz.nx alias.

To simplify slicing named axes, Penzai also provides a helper object:

v1.pz.slice

Builds a slice when sliced (e.g. pz.slice[1:3] == slice(1, 3, None)).

Shape-Checking

pz.chk is an alias for penzai.core.shapecheck, which contains utilities for checking the shapes of PyTrees of positional and named arrays. Some commonly-used attributes on pz.chk:

v1.pz.chk.ArraySpec	Alias of penzai.core.shapecheck.ArraySpec: A non-leaf marker for a (named) array structure.
v1.pz.chk.var(name)	Alias of penzai.core.shapecheck.var: Creates a variable for an axis shape.
v1.pz.chk.vars_for_axes(var_name, ...)	Alias of penzai.core.shapecheck.vars_for_axes: Creates variables for a known collection of named axes.

See penzai.core.shapecheck for documentation of all of the methods and classes accessible through the pz.chk alias.

Dataclass and Struct Utilities

v1.pz.is_pytree_dataclass_type(cls)	Alias of penzai.core.struct.is_pytree_dataclass_type: Checks if a class was wrapped in the pytree_dataclass decorator.
v1.pz.is_pytree_node_field(field)	Alias of penzai.core.struct.is_pytree_node_field: Returns True if this field is treated as a PyTree child node by Struct.
v1.pz.StructStaticMetadata	Alias of penzai.core.struct.StructStaticMetadata: Container for a struct's static fields.
v1.pz.PyTreeDataclassSafetyError	Alias of penzai.core.struct.PyTreeDataclassSafetyError: Error raised due to pytree dataclass safety checks.

Rendering and Global Configuration Management

These utilities are available in the pz namespace for backwards compatibility. However, they have been moved to the separate Treescope pretty-printing package. See the Treescope documentation for more information.

v1.pz.ts	Alias of penzai.pz.ts: Common aliases for treescope.
v1.pz.show(*args[, wrap, space_separated])	Alias of penzai.treescope._compatibility_setup.show: Shows a list of objects inline, like python print, but with rich display.
v1.pz.ContextualValue	Alias of treescope.context.ContextualValue: A global value which can be modified in a scoped context.
v1.pz.oklch_color(lightness, chroma, hue[, ...])	Alias of treescope.formatting_util.oklch_color: Constructs an OKLCH CSS color.
v1.pz.color_from_string(key_string[, ...])	Alias of treescope.formatting_util.color_from_string: Derives a color whose hue is keyed by a string.
v1.pz.dataclass_from_attributes(cls, ...)	Alias of treescope.dataclass_util.dataclass_from_attributes: Directly instantiates a dataclass given all of its fields.
v1.pz.init_takes_fields(cls)	Alias of treescope.dataclass_util.init_takes_fields: Returns True if cls.__init__ takes exactly one argument per field.