import string


def _expand_wildcard(expr: str, num_wildcard_dims: int) -> str:
  """Replaces a wildcard with the correct number of axes.

  Args:
    expr: Shape expression containing '...'.
    num_wildcard_dims: Number of dimensions captured by the wildcard.

  Returns:
    Copy of `expr` with the wildcard '...' replaced by indices 'A', 'B', ...
  """
  # Introduce indices A, B, C, ... for the wildcard.
  if num_wildcard_dims > len(string.ascii_uppercase):
    raise ValueError('Too many dimensions')
  expanded_wildcard = string.ascii_uppercase[:num_wildcard_dims]
  return expr.replace('...', expanded_wildcard)