code: def _check_tree(
      func_name,
      expected_name,
      actual_tree,
      expected_tree,
      has_aux=False,
  ):
      if has_aux:
          actual_tree_children = actual_tree.children()
  
          if len(actual_tree_children) == 2:
              # select first child as result tree
              actual_tree = actual_tree_children[0]
          else:
              raise ValueError(
                  f"{func_name}() produced a pytree with structure "
                  f"{actual_tree}, but a pytree tuple with auxiliary "
                  f"output was expected because has_aux was set to True."
              )
  
      if actual_tree != expected_tree:
          raise TypeError(
              f"{func_name}() output pytree structure must match {expected_name}, "
              f"got {actual_tree} and {expected_tree}."
          )
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import os
  from functools import partial
  from typing import Any, Callable
  from jax._src import core
  from jax._src import linear_util
  from jax._src.lax import lax
  from jax._src import effects
  from jax._src import ad_util
  from jax._src import state
  from jax._src import util
  from jax._src.util import cache, weakref_lru_cache, safe_map, unzip3, partition_list
  from jax.api_util import flatten_fun_nokwargs
  from jax._src.interpreters import partial_eval
  from jax.tree_util import tree_map, tree_unflatten
label_desc: Use image node_id 7 for calling a global function with example usage: _check_tree(func_name, expected_name, actual_tree, expected_tree, has_aux) without return types