def _match_leaf(
    array_structure: ArraySpec,
    value: Any,
) -> list[_MatchConstraint]:
  """Helper function to match a leaf value against an `ArraySpec`.

  Args:
    array_structure: Expected structure to match agains.
    value: Value to match.

  Returns:
    List of constraints that must be satisfied for this value to match the
    structure.
  """
  concrete_structure = _abstract_leaf(value)
  if isinstance(concrete_structure, ArraySpec):
    positional_shape = concrete_structure.shape
    named_shape = concrete_structure.named_shape
    dtype = concrete_structure.dtype
  else:
    return [
        _UnsatisfiedConstraint(
            f"Couldn't match an ArraySpec with a non-arraylike value {value}"
        )
    ]

  constraints = []

  # Check dtype.
  if isinstance(array_structure.dtype, type) and issubclass(
      array_structure.dtype, np.generic
  ):
    if not jnp.issubdtype(dtype, array_structure.dtype):
      constraints.append(
          _UnsatisfiedConstraint(
              "Value has the wrong dtype: expected a sub-dtype of"
              f" {array_structure.dtype} but got dtype {dtype}."
          )
      )
  else:
    expected_dtype = jax.dtypes.canonicalize_dtype(
        array_structure.dtype, allow_extended_dtype=True
    )
    actual_dtype = jax.dtypes.canonicalize_dtype(
        dtype, allow_extended_dtype=True
    )
    if actual_dtype != expected_dtype:
      constraints.append(
          _UnsatisfiedConstraint(
              "Value has the wrong dtype: expected dtype"
              f" {expected_dtype} but got dtype {actual_dtype}."
          )
      )

  constraints.append(
      _PositionalConstraint(
          tuple(positional_shape), tuple(array_structure.shape)
      )
  )
  constraints.append(
      _NamedConstraint(dict(named_shape), dict(array_structure.named_shape))
  )
  return constraints
Helper function to match a leaf value against an `ArraySpec`.

  Args:
    array_structure: Expected structure to match agains.
    value: Value to match.

  Returns:
    List of constraints that must be satisfied for this value to match the
    structure.
  
penzai/core/shapecheck.py