def _try_match_one(
    keypath,
    pattern: int | DimVar | MultiDimVar | KnownDim,
    value: int | tuple[int, ...] | dict[named_axes.AxisName, int],
    solutions: dict[str | tuple[str, str], _Binding],
) -> str | None:
  """Internal helper to match a pattern with a value.

  Args:
    keypath: Keypath for the current value.
    pattern: Pattern to match.
    value: Value to match.
    solutions: Current solutions for variable bindings. New variable values will
      be added to this.

  Returns:
    None if the match succeeded, otherwise a string describing the failure.
  """
  if isinstance(pattern, DimVar | MultiDimVar):
    if isinstance(pattern, DimVar):
      assert isinstance(value, int)
    else:
      assert isinstance(value, (tuple, dict))
    if pattern.name in solutions:
      if solutions[pattern.name].value != value:
        return (
            f"Size {value} does not match previous size"
            f" {solutions[pattern.name].value} for {_show_dim(pattern)} from"
            f" {solutions[pattern.name].source}"
        )
    else:
      solutions[pattern.name] = _Binding(
          "root" + jax.tree_util.keystr(keypath), value
      )
  elif isinstance(pattern, KnownDim):
    assert isinstance(value, int)
    if pattern.name in solutions:
      if solutions[pattern.name].value != pattern.size:
        return (
            f"Size {pattern.size} does not match previous size"
            f" {solutions[pattern.name].value} for {_show_dim(pattern)} from"
            f" {solutions[pattern.name].source}"
        )
    else:
      solutions[pattern.name] = _Binding(
          "root" + jax.tree_util.keystr(keypath), value
      )
    if value != pattern.size:
      if pattern.from_keypath is not None:
        return (
            f"Size {value} does not match previous"
            f" size {pattern.size} for {_show_dim(pattern)} from"
            f" {pattern.from_keypath}"
        )
      else:
        return (
            f"Actual size {value} does not match expected {_show_dim(pattern)}"
        )
  elif isinstance(pattern, int):
    if value != pattern:
      return f"Actual size {value} does not match expected {_show_dim(pattern)}"
  return None
Internal helper to match a pattern with a value.

  Args:
    keypath: Keypath for the current value.
    pattern: Pattern to match.
    value: Value to match.
    solutions: Current solutions for variable bindings. New variable values will
      be added to this.

  Returns:
    None if the match succeeded, otherwise a string describing the failure.
  
penzai/core/shapecheck.py