def check_structure(
    value: Any,
    pattern: StructureAnnotation,
    known_vars: (
        Mapping[str, int | tuple[int, ...] | dict[named_axes.AxisName, int]]
        | None
    ) = None,
    error_prefix: str = "",
) -> MatchResult:
  """Checks that a structure of values matches a pattern.

  Args:
    value: A PyTree of arrays to check.
    pattern: A PyTree of `ArraySpec` leaves to check against. Each `ArraySpec`
      can include dimension variables (created by ``var(name)``, ``*var(name)``,
      ``**var(name)`` or ``**vars_for_axes(name, spec)``), and any two variables
      with the same name must have the same length. The pattern may also include
      other leaves which will be checked for exact equality.
    known_vars: A previous match result (or other assignment of names to values)
      that the names must be consistent with.
    error_prefix: Optional prefix to prepend to exceptions to indicate the
      context in which shape-checking is being done.

  Returns:
    A match result containing the values for each dimension variable in the
    pattern.

  Raises:
    StructureMismatchError: If the value does not match the pattern.
  """
  all_constraints = []

  def add_constraints(keypath, pattern: Any, value: Any):
    if isinstance(pattern, ArraySpec):
      for constraint in _match_leaf(pattern, value):
        all_constraints.append((keypath, constraint))
    elif isinstance(pattern, Wildcard):
      # Wildcard matches anything.
      pass
    else:
      is_equal = pattern == value
      if isinstance(is_equal, jax.Array | np.ndarray):
        try:
          is_equal = bool(np.all(is_equal))
        except jax.errors.ConcretizationTypeError:
          is_equal = False
      if not is_equal:
        failed = _UnsatisfiedConstraint(
            f"Value {repr(value)} was not equal to the non-ArraySpec"
            f" pattern {repr(pattern)}."
        )
        all_constraints.append((keypath, failed))

  jax.tree_util.tree_map_with_path(
      add_constraints,
      pattern,
      value,
      is_leaf=lambda x: isinstance(x, ArraySpec | Wildcard),
  )

  # Phase 1: Collect and handle constraints that don't require solving for
  # variables.
  failures = []
  variable_constraints = []
  for keypath, constraint in all_constraints:
    if isinstance(constraint, _UnsatisfiedConstraint):
      failures.append((keypath, constraint.msg))
    elif isinstance(constraint, _PositionalConstraint):
      variable_constraints.append((keypath, constraint))
    elif isinstance(constraint, _NamedConstraint):
      variable_constraints.append((keypath, constraint))
    else:
      raise ValueError(
          f"Ungiven constraint type {type(constraint)}: {constraint}"
      )

  # Phase 2: Solve for variable constraints.
  # Solutions maps variable names to _Binding instances.
  if known_vars is None:
    solutions = {}
  else:
    solutions = {
        name: _Binding(
            source=(
                "the known variable assignments (argument `known_vars` to"
                " check_structure)"
            ),
            value=value,
        )
        for name, value in dict(known_vars).items()
    }
  while variable_constraints or failures:
    # Try to solve each constraint in order.
    unsolved_variable_constraints = []
    remaining_variables = set()
    for keypath, constraint in variable_constraints:
      if isinstance(constraint, _PositionalConstraint):
        # Solving a positional constraint.
        # First inline any "unpack" patterns that we already have solved:
        constraint, inline_info = _positional_inline_multidimvars(
            constraint, solutions
        )
        if constraint.num_multi_dim_vars() >= 2:
          # Can't solve with multiple dim vars!
          unsolved_variable_constraints.append((keypath, constraint))
          remaining_variables.update(constraint.extract_vars())
        else:
          # We can solve this positional constraint.
          multi_dim_set = [
              i
              for i, val in enumerate(constraint.pattern)
              if isinstance(val, MultiDimVar)
          ]
          if multi_dim_set:
            # We have a single unpacked variable that should match many dims.
            # We need to match prefixes and suffixes, and assign the rest to
            # the wildcard.
            [prefix_size] = multi_dim_set
            suffix_size = len(constraint.pattern) - prefix_size - 1
            if len(constraint.value) < prefix_size + suffix_size:
              failures.append((
                  keypath,
                  (
                      f"Positional shape {constraint.value} was shorter than"
                      " wildcard pattern"
                      f" {_summarize_pattern(constraint.pattern)}"
                      + inline_info
                  ),
              ))
            else:
              pattern_prefix = constraint.pattern[:prefix_size]
              pattern_suffix = constraint.pattern[prefix_size + 1 :]
              multi_dim = constraint.pattern[prefix_size]
              shape_prefix = constraint.value[:prefix_size]
              shape_suffix = constraint.value[
                  len(constraint.value) - suffix_size :
              ]
              shape_middle = constraint.value[
                  prefix_size : len(constraint.value) - suffix_size
              ]

              subfailures = []
              err = _try_match_one(
                  keypath=keypath,
                  pattern=multi_dim,
                  value=shape_middle,
                  solutions=solutions,
              )
              if err is not None:
                subfailures.append(
                    f"  Unpack pattern *var({multi_dim.name}): " + err
                )
              matched_parts = itertools.chain(
                  zip(range(prefix_size), shape_prefix, pattern_prefix),
                  zip(
                      range(
                          len(constraint.value) - suffix_size,
                          len(constraint.value),
                      ),
                      shape_suffix,
                      pattern_suffix,
                  ),
              )
              for i, sval, spat in matched_parts:
                err = _try_match_one(
                    keypath=keypath,
                    pattern=spat,
                    value=sval,
                    solutions=solutions,
                )
                if err is not None:
                  subfailures.append(f"  Dim {i}: " + err)
              if subfailures:
                failures.append((
                    keypath,
                    (
                        "Positional shape mismatch between value"
                        f" {constraint.value} and pattern"
                        f" {_summarize_pattern(constraint.pattern)}:\n"
                        + "\n".join(subfailures)
                        + inline_info
                    ),
                ))
          else:
            # Matching two sequences without dimension wildcards.
            if len(constraint.value) != len(constraint.pattern):
              failures.append((
                  keypath,
                  (
                      f"Positional shape {constraint.value} had different"
                      " length than pattern"
                      f" {_summarize_pattern(constraint.pattern)}"
                      + inline_info
                  ),
              ))
            else:
              subfailures = []
              for i, (sval, spat) in enumerate(
                  zip(constraint.value, constraint.pattern)
              ):
                err = _try_match_one(
                    keypath=keypath,
                    pattern=spat,
                    value=sval,
                    solutions=solutions,
                )
                if err is not None:
                  subfailures.append(f"  Dim {i}: " + err)
              if subfailures:
                failures.append((
                    keypath,
                    (
                        "Positional shape mismatch between value"
                        f" {constraint.value} and pattern"
                        f" {_summarize_pattern(constraint.pattern)}:\n"
                        + "\n".join(subfailures)
                        + inline_info
                    ),
                ))
      elif isinstance(constraint, _NamedConstraint):
        # Solving a named axis constraint.
        # First inline any "unpack" patterns that we already have solved:
        constraint, inline_info = _named_inline_multidimvars(
            constraint, solutions
        )
        if isinstance(constraint, _UnsatisfiedConstraint):
          failures.append((keypath, constraint.msg + inline_info))
        elif constraint.num_multi_dim_vars() >= 2:
          # Can't solve with multiple dim vars!
          unsolved_variable_constraints.append((keypath, constraint))
          remaining_variables.update(constraint.extract_vars())
        else:
          # Should be able to solve.
          non_dimvar_keys = [
              key
              for key in constraint.pattern.keys()
              if not isinstance(key, MultiDimVar)
          ]
          maybe_dimvar_key = [
              key
              for key in constraint.pattern.keys()
              if isinstance(key, MultiDimVar)
          ]
          remaining_value_keys = ordered_set.OrderedSet(constraint.value.keys())
          subfailures = []
          for key in non_dimvar_keys:
            # Check one key.
            if key in constraint.value:
              err = _try_match_one(
                  keypath=keypath,
                  pattern=constraint.pattern[key],
                  value=constraint.value[key],
                  solutions=solutions,
              )
              if err is not None:
                subfailures.append(f"  Axis {repr(key)}: " + err)
              remaining_value_keys.remove(key)
            else:
              subfailures.append(
                  f"  Axis {repr(key)}: Expected to be present but was missing"
              )
          if maybe_dimvar_key:
            # Match the unpack pattern.
            [dimvar_key] = maybe_dimvar_key
            err = _try_match_one(
                keypath=keypath,
                pattern=dimvar_key,
                value={k: constraint.value[k] for k in remaining_value_keys},
                solutions=solutions,
            )
            if err is not None:
              subfailures.append(
                  f"  Unpack pattern **{repr(dimvar_key)}: " + err
              )
          elif remaining_value_keys:
            # No unpack pattern, but we ran out of keys.
            subfailures.append(
                "  Unexpected names in value's named shape:"
                f" {sorted(remaining_value_keys)}"
            )
          if subfailures:
            failures.append((
                keypath,
                (
                    "Named shape mismatch between value"
                    f" {constraint.value} and pattern"
                    f" {_summarize_pattern(constraint.pattern)}:\n"
                    + "\n".join(subfailures)
                    + inline_info
                ),
            ))
      else:
        raise ValueError(
            f"Unknown constraint type {type(constraint)}: {constraint}"
        )

    # If we encountered an error, or if we were unable to resolve any of the
    # constraints, the match is a failure.
    if failures:
      msgs = [error_prefix + "Mismatch while checking structures:"]
      for keypath, failure in failures:
        msgs.append(f"At root{jax.tree_util.keystr(keypath)}: {failure}")
      if unsolved_variable_constraints:
        msgs.append(
            "There may also be additional mismatches due to remaining"
            " unsolved constraints."
        )
      raise StructureMismatchError("\n".join(msgs))
    elif len(unsolved_variable_constraints) < len(variable_constraints):
      variable_constraints = unsolved_variable_constraints
    else:
      raise StructureMismatchError(
          error_prefix
          + "Could not solve for all variable constraints. This usually means"
          " the variable assignment is ambiguous because multiple unpack"
          " patterns (*var(...) or **var(...)) appeared in the same positional"
          " or named shape.\nUnsolvable variables: "
          + ", ".join(_show_dim(var) for var in sorted(remaining_variables))
          + "\nUnsolvable constraints:"
          + "".join(
              f"\n  {constraint.value} =="
              f" {_summarize_pattern(constraint.pattern)} from"
              f" root{jax.tree_util.keystr(keypath)}"
              for (keypath, constraint) in unsolved_variable_constraints
          )
      )

  # No more constraints, which means this was a successful match. We need to
  # re-associate any variables whose values were tuples.
  final_matches = {}
  match_errors = []
  for name, binding in solutions.items():
    if isinstance(name, tuple):
      if name[0] in solutions:
        # Redundant! We already have this in our solution.
        found = None
        if isinstance(name[1], int):
          if isinstance(solutions[name[0]].value, tuple) and name[1] < len(
              solutions[name[0]].value
          ):
            found = solutions[name[0]].value[name[1]]
        else:
          assert isinstance(name[1], str)
          if isinstance(solutions[name[0]].value, dict):
            found = solutions[name[0]].value.get(name[1])
        if found != binding.value:
          match_errors.append(
              f"  {_show_dim(DimVar(name[0]))}: Solved as"
              f" {solutions[name[0]].value} from"
              f" {solutions[name[0]].source} but"
              f" separately solved {_show_dim(DimVar(name))} as"
              f" {binding.value} from"
              f" {binding.source}"
          )
      else:
        if name[0] not in final_matches:
          final_matches[name[0]] = {}
        # Each match should appear in the solution dict at most once, because
        # if it already exists we just check for it.
        assert name[1] not in final_matches[name[0]]
        final_matches[name[0]][name[1]] = binding.value
    else:
      final_matches[name] = binding.value

  if match_errors:
    raise StructureMismatchError(
        error_prefix
        + "Unexpected conflicts between variable assignments:\n"
        + "\n".join(match_errors)
    )
  return MatchResult(final_matches)
Checks that a structure of values matches a pattern.

  Args:
    value: A PyTree of arrays to check.
    pattern: A PyTree of `ArraySpec` leaves to check against. Each `ArraySpec`
      can include dimension variables (created by ``var(name)``, ``*var(name)``,
      ``**var(name)`` or ``**vars_for_axes(name, spec)``), and any two variables
      with the same name must have the same length. The pattern may also include
      other leaves which will be checked for exact equality.
    known_vars: A previous match result (or other assignment of names to values)
      that the names must be consistent with.
    error_prefix: Optional prefix to prepend to exceptions to indicate the
      context in which shape-checking is being done.

  Returns:
    A match result containing the values for each dimension variable in the
    pattern.

  Raises:
    StructureMismatchError: If the value does not match the pattern.
  
penzai/core/shapecheck.py