def as_array_structure(tree: Any) -> StructureAnnotation:
  """Abstracts a concrete tree of values into a tree of `ArraySpec` nodes.

  Args:
    tree: The tree to get an array structure for.

  Returns:
    A copy of `tree` where each arraylike or `NamedArray` subtree is replaced
    by an `ArraySpec`, and each other leaf is replaced by `ANY`.
  """
  return jax.tree_util.tree_map(
      _abstract_leaf, tree, is_leaf=named_axes.is_namedarray
  )
Abstracts a concrete tree of values into a tree of `ArraySpec` nodes.

  Args:
    tree: The tree to get an array structure for.

  Returns:
    A copy of `tree` where each arraylike or `NamedArray` subtree is replaced
    by an `ArraySpec`, and each other leaf is replaced by `ANY`.
  
penzai/core/shapecheck.py