pz.nn: Neural network alias namespace

Parameters

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

Basic Combinators

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

Basic Operations

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

Linear and Affine Layers

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

Standardization

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

Dropout

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

Language Modeling

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