shapecheck

Utilities for runtime shape-checking.

Penzai is highly dynamic by nature, and is heavily based on having nested tree structured objects, which is difficult to capture with Python’s type annotations. This module provides a shape-and-structure checking utility that can be used to check the shapes and dtypes of nested structures while also simultaneously extracting parts that the program may need to use (e.g. lengths for axes).

Classes

ArraySpec	A non-leaf marker for a (named) array structure.
DimVar	A dimension variable, which will be matched with the shapes in a structure.
DimensionVariableSubstitution	A substitution for all of the dimension variables in a structure.
KnownDim	A dimension with a known value, used to bind a name to a value directly.
MatchResult	The result of a successful shape check.
MultiDimVar	A variable standing in for a sequence or dictionary axis names.
RemainingAxisPlaceholder	Sentinel marker for the remainder of axes.
Wildcard	Wildcard marker that matches any value or subtree of a PyTree.

Functions

as_array_structure(tree)	Abstracts a concrete tree of values into a tree of ArraySpec nodes.
check_structure(value, pattern[, ...])	Checks that a structure of values matches a pattern.
full_substitute_dimension_variables(...)	Substitutes all dimension variables for their known values.
get_dimension_variables(structure)	Returns a collection of all dimension variables in a structure.
var(name)	Creates a variable for an axis shape.
vars_for_axes(var_name, axis_names_or_specs)	Creates variables for a known collection of named axes.

Exceptions

StructureMismatchError

Raised when a check_structure call fails.