code: def _check_shapes(func_name, expected_name, actual, expected):
      actual_shapes = _map(np.shape, tree_leaves(actual))
      expected_shapes = _map(np.shape, tree_leaves(expected))
      if actual_shapes != expected_shapes:
          raise ValueError(
              f"{func_name}() output shapes must match {expected_name}, "
              f"got {actual_shapes} and {expected_shapes}"
          )
  
imports: import collections
  from functools import partial
  import operator
  import jax
  from jax.tree_util import tree_flatten, treedef_children, tree_leaves, tree_unflatten, treedef_tuple
  from jax._src import ad_util
  from jax._src import core
  from jax._src import linear_util
  from jax._src.core import raise_to_shaped
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters import mlir
  from jax._src.interpreters import xla
  from jax._src.traceback_util import api_boundary
  from jax._src.util import split_list, safe_map
  import numpy
  from jax._src.lax.control_flow.common import _abstractify, _check_tree, _initial_style_jaxpr
label_desc: Use image node_id 8 for calling a global function with example usage: _check_shapes(func_name, expected_name, actual, expected) without return types