def vars_for_axes(
    var_name: str,
    axis_names_or_specs: Collection[str] | Mapping[str, int | None],
) -> dict[str, DimVar | KnownDim]:
  """Creates variables for a known collection of named axes.

  Args:
    var_name: A name for the variable that will store the concrete sizes for all
      of these named axes.
    axis_names_or_specs: Either a collection of axis names, or a mapping from
      axis names to integer axis sizes or to None (for unknown axis sizes).

  Returns:
    A dictionary with the keys from ``axis_names_or_specs``, and values that
    either reflect either the known size from the value of
    ``axis_names_or_specs`` or an unknown dimension. Intended to be passed as
    the ``named_shape`` of an `ArraySpec` or unpacked into a larger dictionary
    of named shapes.
  """
  if not isinstance(axis_names_or_specs, Mapping):
    axis_names_or_specs = {k: None for k in axis_names_or_specs}

  return {
      k: (
          DimVar((var_name, k))
          if size is None
          else KnownDim((var_name, k), size)
      )
      for k, size in axis_names_or_specs.items()
  }
Creates variables for a known collection of named axes.

  Args:
    var_name: A name for the variable that will store the concrete sizes for all
      of these named axes.
    axis_names_or_specs: Either a collection of axis names, or a mapping from
      axis names to integer axis sizes or to None (for unknown axis sizes).

  Returns:
    A dictionary with the keys from ``axis_names_or_specs``, and values that
    either reflect either the known size from the value of
    ``axis_names_or_specs`` or an unknown dimension. Intended to be passed as
    the ``named_shape`` of an `ArraySpec` or unpacked into a larger dictionary
    of named shapes.
  
penzai/core/shapecheck.py