check_structure

check_structure#

penzai.core.shapecheck.check_structure(value: Any, pattern: StructureAnnotation, known_vars: Mapping[str, int | tuple[int, ...] | dict[named_axes.AxisName, int]] | None = None, error_prefix: str = '') → MatchResult[source]#

Checks that a structure of values matches a pattern.

Parameters:

  • value – A PyTree of arrays to check.

  • pattern – A PyTree of ArraySpec leaves to check against. Each ArraySpec can include dimension variables (created by var(name), *var(name), **var(name) or **vars_for_axes(name, spec)), and any two variables with the same name must have the same length. The pattern may also include other leaves which will be checked for exact equality.

  • known_vars – A previous match result (or other assignment of names to values) that the names must be consistent with.

  • error_prefix – Optional prefix to prepend to exceptions to indicate the context in which shape-checking is being done.

Returns:

A match result containing the values for each dimension variable in the pattern.

Raises:

StructureMismatchError – If the value does not match the pattern.

Contents