pz: Penzai’s alias namespace

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. Most models and layers are subclasses of pz.Layer, which means they can be called with a single argument.

pz.pytree_dataclass([cls, ...])	Alias of penzai.core.struct.pytree_dataclass: Decorator for constructing a frozen PyTree dataclass.
pz.Struct	Alias of penzai.core.struct.Struct: Base class for penzai PyTree structures.
pz.Layer	Alias of penzai.core.layer.Layer: Abstract base class for neural network layers and other 1-arg callables.
pz.LayerLike	alias of Callable[[Any], Any]

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.

pz.select(tree)	Alias of penzai.core.selectors.select: Wraps a tree in a singleton selection for processing.
pz.Selection	Alias of penzai.core.selectors.Selection: A selected subset of nodes within a larger PyTree.
pz.combine(*partitions)	Alias of penzai.core.partitioning.combine: Combines leaves from multiple partitions.
pz.NotInThisPartition	Alias of penzai.core.partitioning.NotInThisPartition: Sentinel object that identifies subtrees removed by partition.
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:

pz.nx.NamedArray	Alias of penzai.core.named_axes.NamedArray: A multidimensional array with a combination of positional and named axes.
pz.nx.nmap(fun)	Alias of penzai.core.named_axes.nmap: Automatically vectorizes fun over named axes of NamedArray inputs.
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:

pz.slice

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

Visualization

pz.ts is an alias namespace for Penzai’s interactive pretty printer Treescope. Some commonly-used attributes on pz.ts:

pz.ts.register_as_default([...])	Alias of penzai.treescope.treescope_ipython.register_as_default: Registers treescope as the default IPython renderer.
pz.ts.register_autovisualize_magic()	Alias of penzai.treescope.treescope_ipython.register_autovisualize_magic: Registers the %%autovisualize magic.
pz.ts.render_array(array, *[, columns, ...])	Alias of penzai.treescope.arrayviz.arrayviz.render_array: Renders an array (positional or named) to a displayable HTML object.

See the documentation for pz.ts to view all of the methods and classes accessible through this alias namespace.

Penzai also provides a utility for quickly showing a value with Treescope in an IPython notebook, using syntax similar to ordinary print:

pz.show(*args[, wrap, space_separated])

Alias of penzai.treescope.treescope_ipython.show: Shows a list of objects inline, like python print, but with rich display.

Neural Networks

pz.nn is an alias namespace for Penzai’s declarative neural network system, which uses a combinator-based design to expose all of your model’s operations as nodes in your model PyTree. pz.nn re-exports layers from submodules of penzai.nn in a single convenient namespace.

See the documentation for pz.nn to view all of the methods and classes accessible through this alias namespace.

Data Effects

pz.de is an alias namespace for Penzai’s “data effect” system, which represents side inputs, side outputs, randomness, and mutable state as typed attributes inside your model PyTree, and allows handling them using functional effect handlers. pz.de re-exports effect requests and effect handlers from submodules of penzai.data_effects in a single convenient namespace.

See the documentation for pz.de to view all of the methods and classes accessible through this alias namespace.

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:

pz.chk.ArraySpec	Alias of penzai.core.shapecheck.ArraySpec: A non-leaf marker for a (named) array structure.
pz.chk.var(name)	Alias of penzai.core.shapecheck.var: Creates a variable for an axis shape.
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.

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

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

Context Management

pz.disable_interactive_context()	Alias of penzai.core.context.disable_interactive_context: Clears the global interactive context stack and disables interactive mode.
pz.enable_interactive_context()	Alias of penzai.core.context.enable_interactive_context: Enables the global interactive context stack.
pz.ContextualValue	Alias of penzai.core.context.ContextualValue: A global value which can be modified in a scoped context.

Dataclass and Struct Utilities

pz.dataclass_from_attributes(cls, **field_values)	Alias of penzai.core.dataclass_util.dataclass_from_attributes: Directly instantiates a dataclass given all of its fields.
pz.init_takes_fields(cls)	Alias of penzai.core.dataclass_util.init_takes_fields: Returns True if cls.__init__ takes exactly one argument per field.
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.
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.
pz.StructStaticMetadata	Alias of penzai.core.struct.StructStaticMetadata: Container for a struct's static fields.
pz.PyTreeDataclassSafetyError	Alias of penzai.core.struct.PyTreeDataclassSafetyError: Error raised due to pytree dataclass safety checks.

Context and State Management

pz.disable_interactive_context()	Alias of penzai.core.context.disable_interactive_context: Clears the global interactive context stack and disables interactive mode.
pz.enable_interactive_context()	Alias of penzai.core.context.enable_interactive_context: Enables the global interactive context stack.
pz.ContextualValue	Alias of penzai.core.context.ContextualValue: A global value which can be modified in a scoped context.
pz.RandomStream	Alias of penzai.core.random_stream.RandomStream: Helper object to construct a stream of random numbers.

Formatting Utilities

pz.oklch_color(lightness, chroma, hue[, alpha])	Alias of penzai.core.formatting_util.oklch_color: Constructs an OKLCH CSS color.
pz.color_from_string(key_string[, ...])	Alias of penzai.core.formatting_util.color_from_string: Derives a color whose hue is keyed by a string.