def full_substitute_dimension_variables(
    structure: StructureAnnotation,
    substitution: DimensionVariableSubstitution,
) -> Any:
  """Substitutes all dimension variables for their known values.

  Args:
    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).
  """

  def subst_leaf(keypath, leaf):
    if isinstance(leaf, ArraySpec):
      new_shape = []
      for dim in leaf.shape:
        if isinstance(dim, DimVar):
          assignment = substitution.size_variables[dim.name]
          new_shape.append(assignment)
        elif isinstance(dim, KnownDim):
          new_shape.append(dim.size)
        elif isinstance(dim, MultiDimVar):
          assignment = substitution.sequence_variables[dim.name]
          new_shape.extend(assignment)
        else:
          new_shape.append(dim)
      new_named_shape = {}
      for key, val in leaf.named_shape.items():
        if isinstance(val, DimVar):
          assignment = substitution.size_variables[val.name]
          new_named_shape[key] = assignment
        elif isinstance(val, KnownDim):
          new_named_shape[key] = val.size
        elif isinstance(key, MultiDimVar):
          assert isinstance(val, RemainingAxisPlaceholder)
          assignment = substitution.mapping_variables[key.name]
          for subkey, subval in assignment.items():
            if subkey in new_named_shape or subkey in leaf.named_shape:
              raise ValueError(
                  "Key conflict while substituting mapping variable"
                  f" {_show_dim(key)} at root{jax.tree_util.keystr(keypath)}:"
                  f" name {subkey} is already present"
              )
            new_named_shape[subkey] = subval
        else:
          new_named_shape[key] = val
      return ArraySpec(
          shape=tuple(new_shape), named_shape=new_named_shape, dtype=leaf.dtype
      )
    else:
      return leaf

  return jax.tree_util.tree_map_with_path(
      subst_leaf, structure, is_leaf=lambda x: isinstance(x, ArraySpec)
  )
Substitutes all dimension variables for their known values.

  Args:
    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).
  
penzai/core/shapecheck.py