code: def batch(
      kernel_fn: _KernelFn,
      batch_size: int = 0,
      device_count: int = -1,
      store_on_device: bool = True,
  ) -> _KernelFn:
      """Returns a function that computes a kernel in batches over all devices.
  
      Note that you typically should not apply the `jax.jit` decorator to the
      resulting `batched_kernel_fn`, as its purpose is explicitly serial execution
      in order to save memory. Further, you do not need to apply `jax.jit` to the
      input `kernel_fn` function, as it is JITted internally.
  
      Args:
        kernel_fn:
          A function that computes a kernel on two batches,
          `kernel_fn(x1, x2, *args, **kwargs)`. Here `x1` and `x2` are
          `jnp.ndarray`s of shapes `(n1,) + input_shape` and `(n2,) + input_shape`.
          The kernel function should return a `PyTree`.
  
        batch_size:
          specifies the size of each batch that gets processed per physical device.
          Because we parallelize the computation over columns it should be the case
          that `x1.shape[0]` is divisible by `device_count * batch_size` and
          `x2.shape[0]` is divisible by `batch_size`.
  
        device_count:
          specifies the number of physical devices to be used. If
          `device_count == -1` all devices are used. If `device_count == 0`, no
          device parallelism is used (a single default device is used).
  
        store_on_device:
          specifies whether the output should be kept on device or brought back to
          CPU RAM as it is computed. Defaults to `True`. Set to `False` to store
          and concatenate results using CPU RAM, allowing to compute larger kernels.
  
      Returns:
        A new function with the same signature as `kernel_fn` that computes the
        kernel by batching over the dataset in parallel with the specified
        `batch_size` using `device_count` devices.
      """
      # TODO(romann): find a way to avoid reading requirements.
      input_req = getattr(kernel_fn, "input_req", {})
      dropout_in_analytic_kernel = input_req.get("use_dropout", False)
      use_multidevice = device_count > 0 or (
          device_count == -1 and jax.local_device_count() > 1
      )
      use_serial = bool(batch_size)
      if use_multidevice:
          kernel_fn = _parallel(
              kernel_fn,
              use_serial,
              dropout_in_analytic_kernel,
              device_count,
          )
      else:
          kernel_fn = _jit_or_pmap_broadcast(kernel_fn, 0)
  
      if not use_serial:
          return kernel_fn
  
      return _serial(kernel_fn, batch_size, store_on_device)
  
imports: from functools import partial
  from typing import Any, Callable, Iterable, Optional, TypeVar, Union
  import warnings
  import jax
  from jax import device_put
  from jax import devices
  from jax import jit
  from jax import pmap
  from jax import random
  import jax.numpy
  from jax.tree_util import tree_all
  from jax.tree_util import tree_flatten
  from jax.tree_util import tree_map
  from jax.tree_util import tree_unflatten
  import numpy
  from .utils import utils
  from .utils.kernel import Kernel
  from .utils.typing import KernelFn
  from .utils.typing import NTTree
label_desc: Use image node_id 1 for calling a global function with example usage: batch(kernel_fn, batch_size, device_count, store_on_device) and returns: _serial, kernel_fn