def _abstract_leaf(value: Any) -> ArraySpec | Wildcard:
  """Helper function to get an `ArraySpec` view of a leaf."""
  if isinstance(value, named_axes.NamedArrayBase):
    value.check_valid()
    positional_shape = value.positional_shape
    named_shape = value.named_shape
    dtype = value.dtype
  elif isinstance(value, ArraySpec):
    positional_shape = value.shape
    named_shape = value.named_shape
    dtype = value.dtype
  elif hasattr(value, "shape") and hasattr(value, "dtype"):
    positional_shape = value.shape
    named_shape = {}
    dtype = value.dtype
  elif isinstance(value, jax.typing.ArrayLike):
    positional_shape = jnp.shape(value)
    named_shape = {}
    dtype = jnp.result_type(value)
  else:
    return ANY
  return ArraySpec(shape=positional_shape, named_shape=named_shape, dtype=dtype)
Helper function to get an `ArraySpec` view of a leaf.
penzai/core/shapecheck.py