full_substitute_dimension_variables

full_substitute_dimension_variables#

penzai.core.shapecheck.full_substitute_dimension_variables(structure: StructureAnnotation, substitution: DimensionVariableSubstitution) Any[source]#

Substitutes all dimension variables for their known values.

Parameters:
  • structure – A PyTree structure containing ArraySpec leaves.

  • substitution – A substitution for the dimension variables. Should have values for all variables in the structure, usually with the same keys as returned by get_dimension_variables.

Returns:

A copy of the structure with all dimension variables substituted for their known values. Any KnownDim instances will also be stripped out and replaced by their constant sizes, so if the substitution did not contain any new dimension variables, the resulting structure will be fully concrete (with no variables).