pz.nn: Neural network alias namespace

Layers and Parameter Utilities

pz.nn.Layer	Alias of penzai.nn.layer.Layer: Abstract base class for neural network layers and other 1-arg callables.
pz.nn.ParameterLike	Alias of penzai.nn.parameters.ParameterLike: Protocol for a parameter-like object.
pz.nn.derive_param_key(base_rng, name)	Alias of penzai.nn.parameters.derive_param_key: Derives a PRNG key for a parameter from a base key and a name.
pz.nn.make_parameter(name, init_base_rng, ...)	Alias of penzai.nn.parameters.make_parameter: Makes a parameter variable (or slot) with a given name and initializer.
pz.nn.assert_no_parameter_slots(model)	Alias of penzai.nn.parameters.assert_no_parameter_slots: Asserts that the given model has no ParameterSlot subtrees.

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 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.
pz.nn.TanhSoftCap	Alias of penzai.nn.basic_ops.TanhSoftCap: Softly rescales a value to lie within a range using tanh.

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.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.ApplyExplicitAttentionMask	Alias of penzai.nn.attention.ApplyExplicitAttentionMask: Applies an explicit attention mask to its input logit array.
pz.nn.ApplyCausalAttentionMask	Alias of penzai.nn.attention.ApplyCausalAttentionMask: Builds and applies a causal attention mask based on token positions.
pz.nn.ApplyCausalSlidingWindowAttentionMask	Alias of penzai.nn.attention.ApplyCausalSlidingWindowAttentionMask: Builds and applies a sliding-window attention mask using token positions.
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).

Layer Stacks

pz.nn.LayerStack	Alias of penzai.nn.layer_stack.LayerStack: A sequence of layers with identical structure, called under jax.lax.scan.
pz.nn.LayerStackVarBehavior	Alias of penzai.nn.layer_stack.LayerStackVarBehavior: Behavior of a variable in a layer stack.
pz.nn.layerstack_axes_from_keypath(keypath)	Alias of penzai.nn.layer_stack.layerstack_axes_from_keypath: Extracts the stacked axes from a keypath.
pz.nn.LayerStackGetAttrKey	Alias of penzai.nn.layer_stack.LayerStackGetAttrKey: GetAttrKey for LayerStack with extra metadata.