get_dimension_variables

penzai.core.shapecheck.get_dimension_variables(structure: StructureAnnotation) → DimensionVariableSubstitution

Returns a collection of all dimension variables in a structure.

Parameters:

structure – The structure to extract variables from. Usually a PyTree with ArraySpec leaves.

Returns:

A DimensionVariableSubstitution that contains all dimension variables in the structure, mapping them back to themselves, so that calling full_substitute_dimension_variables would leave the result unchanged (except for KnownDim instances). Usually, you will want to replace the values in the substitution with something else, and keep only the keys.