def get_dimension_variables(
    structure: StructureAnnotation,
) -> DimensionVariableSubstitution:
  """Returns a collection of all dimension variables in a structure.

  Args:
    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.
  """
  result = DimensionVariableSubstitution({}, {}, {})

  def get_leaf(leaf):
    if isinstance(leaf, ArraySpec):
      for dim in leaf.shape:
        if isinstance(dim, DimVar):
          result.size_variables[dim.name] = dim
        elif isinstance(dim, MultiDimVar):
          result.sequence_variables[dim.name] = (dim,)
      for key, val in leaf.named_shape.items():
        if isinstance(val, DimVar):
          result.size_variables[val.name] = val
        elif isinstance(key, MultiDimVar):
          assert isinstance(val, RemainingAxisPlaceholder)
          result.mapping_variables[key.name] = {key: val}

  jax.tree_util.tree_map(
      get_leaf, structure, is_leaf=lambda x: isinstance(x, ArraySpec)
  )
  return result
Returns a collection of all dimension variables in a structure.

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