def _positional_inline_multidimvars(
    constraint: _PositionalConstraint,
    solutions: dict[str | tuple[str, str], _Binding],
) -> tuple[_PositionalConstraint, str]:
  """Simplifies a positional constraint by inlining multivars."""
  new_pattern = []
  inline_summary = []
  for dim in constraint.pattern:
    if isinstance(dim, MultiDimVar) and dim.name in solutions:
      binding = solutions[dim.name]
      assert isinstance(binding.value, tuple)
      for i, val in enumerate(binding.value):
        new_pattern.append(
            KnownDim((dim.name, i), val, from_keypath=binding.source)
        )
      inline_summary.append(
          f"\n  After inlining var({repr(dim.name)}) = {binding.value} from"
          f" {binding.source}"
      )
    else:
      new_pattern.append(dim)
  return (
      _PositionalConstraint(value=constraint.value, pattern=tuple(new_pattern)),
      "".join(inline_summary),
  )
Simplifies a positional constraint by inlining multivars.
penzai/core/shapecheck.py