def tree_flatten_exactly_one_level(
    tree: Any,
) -> Optional[tuple[list[tuple[Any, Any]], PyTreeDef]]:
  """Flattens a PyTree exactly one level, or returns None if it's not a PyTree.

  Args:
    tree: Tree to flatten.

  Returns:
    If ``tree`` has any children, returns a tuple ``(children, treedef)`` where
    children is a list of ``(key, child)`` pairs. Otherwise, returns ``None``.
  """
  paths_and_subtrees, treedef = jax.tree_util.tree_flatten_with_path(
      tree, is_leaf=lambda subtree: subtree is not tree
  )
  leaf_treedef = jax.tree_util.tree_structure(1)
  if treedef == leaf_treedef:
    return None

  keys_and_subtrees = [
      (key, subtree) for ((key,), subtree) in paths_and_subtrees
  ]
  return keys_and_subtrees, treedef
Flattens a PyTree exactly one level, or returns None if it's not a PyTree.

  Args:
    tree: Tree to flatten.

  Returns:
    If ``tree`` has any children, returns a tuple ``(children, treedef)`` where
    children is a list of ``(key, child)`` pairs. Otherwise, returns ``None``.
  
penzai/core/tree_util.py