code: def tree_flatten_one_level(
      pytree: PyTree,
  ) -> tuple[list[PyTree], PyTreeDef]:  # pyright: ignore
      """Returns the immediate subnodes of a PyTree node. If called on a leaf node then it
      will return just that leaf.
  
      **Arguments:**
  
      - `pytree`: the PyTree node to flatten by one level.
  
      **Returns:**
  
      As `jax.tree_util.tree_flatten`: a list of leaves and a `PyTreeDef`.
  
      !!! Example
  
          ```python
          x = {"a": 3, "b": (1, 2)}
          eqx.tree_flatten_one_level(x)
          # ([3, (1, 2)], PyTreeDef({'a': *, 'b': *}))
  
          y = 4
          eqx.tree_flatten_one_level(y)
          # ([4], PyTreeDef(*))
          ```
      """
      seen_pytree = False
  
      def is_leaf(node):
          nonlocal seen_pytree
          if node is pytree:
              if seen_pytree:
                  # We expect to only see it once as the root.
                  # This catches for example
                  # ```python
                  # x = []
                  # x.append(x)
                  # tree_subnodes(x)
                  # ```
                  # Note that it intentionally does *not* catch
                  # ```python
                  # x = []
                  # y = []
                  # x.append(y)
                  # y.append(x)
                  # tree_subnodes(x)
                  # ```
                  # as `x` is not an immediate subnode of itself.
                  # If you want to check for that then use `tree_check`.
                  try:
                      type_string = type(pytree).__name__
                  except AttributeError:
                      type_string = "<unknown>"
                  raise ValueError(
                      f"PyTree node of type `{type_string}` is immediately "
                      "self-referential; that is to say it appears within its own PyTree "
                      "structure as an immediate subnode. (For example "
                      "`x = []; x.append(x)`.) This is not allowed."
                  )
              else:
                  seen_pytree = True
              return False
          else:
              return True
  
      return jtu.tree_flatten(pytree, is_leaf=is_leaf)
  
imports: from collections.abc import Callable, Sequence
  from typing import Any, Optional, TYPE_CHECKING, Union
  import jax.numpy
  import jax.tree_util
  import numpy
  from jaxtyping import Array, ArrayLike, Bool, Float, PyTree, PyTreeDef
  from ._custom_types import sentinel
  from ._doc_utils import doc_repr
  from ._filters import is_array
label_desc: Use image node_id 5 for calling a global function with example usage: tree_flatten_one_level(pytree) and returns: jtu, False, True