DimensionVariableSubstitution

class penzai.core.shapecheck.DimensionVariableSubstitution

Bases: object

A substitution for all of the dimension variables in a structure.

DimensionVariableSubstitution can be used to inspect and modify the dimension variables appearing in an unknown structure. It is intended to be used as part of transformations that inspect the input or output structures of layers.

The substitutions are allowed to contain other dimension variables. For instance, you can set mapping_variables["foo"] = {"bar":1, **var("baz")} to map the original unpacked mapping variable “foo” to the new composite mapping. This can be used to e.g. rename variables.

Variables:

• size_variables (dict[str | tuple[str, str | int], int | DimVar]) – A map from dimension variable names to their sizes.

• sequence_variables (dict[str, tuple[int | DimVar | MultiDimVar, ...]]) – A map from sequence variable names to their sequence of sizes.

• mapping_variables (dict[str, dict[Any | MultiDimVar, int | DimVar | RemainingAxisPlaceholder]]) – A map from mapping variable names to their mappings from axis name to sizes.

Methods

__init__(size_variables, sequence_variables, ...)
is_empty()	Returns True if there are no variables.

Attributes

size_variables
sequence_variables
mapping_variables

is_empty() → bool

Returns True if there are no variables.