code: def check_vmap_rule_trees(
      rule, original_out_tree, out_tree, out_batched_tree
  ):
      if out_tree != out_batched_tree:
          raise ValueError(
              "structure of output value and output batching specification returned "
              f"by custom vmap rule ({rule_name(rule)}) do not match.\n"
              f"Output values: {out_tree}\n"
              f"Batching spec: {out_batched_tree}"
          )
      if out_tree != original_out_tree:
          raise ValueError(
              f"structure of output returned by custom vmap rule ({rule_name(rule)}) "
              "does not match that of original custom-vmapped function.\n"
              f"Original output: {original_out_tree}\n"
              f"Rule output: {out_tree}"
          )
  
imports: from __future__ import annotations
  import functools
  import operator
  from typing import Callable
  from jax import lax
  from jax._src import api
  from jax._src import core
  from jax._src import custom_api_util
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src import traceback_util
  from jax._src import tree_util
  from jax._src import util
  from jax._src.api_util import flatten_fun_nokwargs
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters.batching import not_mapped
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax._src.interpreters import xla
  from jax._src.tree_util import tree_flatten, tree_map, tree_structure, tree_unflatten, treedef_tuple
label_desc: Use image node_id 4 for calling a global function with example usage: check_vmap_rule_trees(rule, original_out_tree, out_tree, out_batched_tree) without return types