code: def io_callback_batching_rule(
      args, dims, callback, result_avals, sharding, ordered
  ):
      if ordered:
          raise ValueError("Cannot `vmap` ordered IO callback.")
      return pure_callback_batching_rule(
          args,
          dims,
          callback=callback,
          sharding=sharding,
          vectorized=False,
          result_avals=result_avals,
      )
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import functools
  from typing import Any, Callable
  import numpy
  from jax._src import core
  from jax._src import dispatch
  from jax._src import dtypes
  from jax._src import effects
  from jax._src import sharding_impls
  from jax._src import tree_util
  from jax._src import util
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters import mlir
  from jax._src.lib import xla_client
  from jax._src.lax.control_flow.loops import map
  from jax._src.sharding_impls import SingleDeviceSharding
label_desc: Use image node_id 15 for calling a global function with example usage: io_callback_batching_rule(args, dims, callback, result_avals, sharding, ordered) and returns: pure_callback_batching_rule