struct

Core struct abstraction.

A Struct is a PyTree dataclass with a few extra features.

In penzai, Struct is used as the primary way to define any class that should behave as a JAX pytree, not just neural network modules.

Classes

AbstractStructMetaclass	The metaclass for penzai.Struct and its (possibly abstract) subclasses.
Struct	Base class for penzai PyTree structures.
StructStaticMetadata	Container for a struct's static fields.

Functions

is_pytree_dataclass_type(cls)	Checks if a class was wrapped in the pytree_dataclass decorator.
is_pytree_node_field(field)	Returns True if this field is treated as a PyTree child node by Struct.
pytree_dataclass([cls, ...])	Decorator for constructing a frozen PyTree dataclass.

Exceptions

PyTreeDataclassSafetyError

Error raised due to pytree dataclass safety checks.