code: def _parallel(
      kernel_fn: _KernelFn,
      use_serial: bool = True,
      dropout_in_analytic_kernel: bool = False,
      device_count: int = -1,
  ) -> _KernelFn:
      """Returns a function that computes a kernel in batches in parallel.
  
      When batching in parallel, the data is split over a set number of devices.
      The number of devices must be less than or equal to the number of physical
      devices. Moreover, the dataset size needs to divide the device count.
  
      Given two datasets `x1` and `x2`, parallel splits the kernel calculation over
      devices such that each device computes a batch of rows of shape
      `[|x1| / device_count, |x2|]`.
  
      Args:
        kernel_fn:
          A function that computes a kernel between two datasets,
          `kernel_fn(x1, x2)` or the compositional kernel for an input kernel
          `kernel_fn(kernel_in)`. Here `x1` and `x2` are `jnp.ndarray`s of floats of
          shape `(n1,) + input_shape` and `(n2,) + input_shape`; `kernel_in` is a
          Kernel object. The kernel function should return a `PyTree`.
  
        use_serial:
          Whether `serial` will be called after `_parallel`. The only use case is to
          make sure when `dropout` is used in the analytic/empirical kernel, the
          batch size in each device is square.
  
        dropout_in_analytic_kernel:
          whether `dropout` is used in the analytic kernel. See `use_serial` above
          for the only use case.
  
        device_count:
          Integer specifying the number of devices over which to split the data. If
          `device_count == 0`, the computation is parallelized over all available
          devices.
  
      Returns:
        A new function with the same signature as kernel_fn that computes the kernel
        by batching over the dataset in parallel over a specified number of cores.
      """
  
      if device_count == -1:
          device_count = jax.local_device_count()
  
      def _check_dropout(n1, n2, kwargs):
          dropout_in_empirical_kernel = (
              getattr(kwargs, "rng", None) is not None
          )
          if (
              n1 == n2
              and (
                  dropout_in_empirical_kernel
                  or dropout_in_analytic_kernel
              )
              and not use_serial
          ):
              raise NotImplementedError(
                  "Batching for empirical / analytic kernels with dropout"
                  " is not implemented for non-square batch size. "
                  "Using `serial` (i.e. use a non-zero batch_size in the "
                  "`batch` function.) could enforce square batch size in each device."
              )
  
      def _get_n_per_device(n1):
          _device_count = device_count
  
          n1_per_device, ragged = divmod(n1, device_count)
          if n1_per_device and ragged:
              raise ValueError(
                  (
                      "Dataset size ({}) must divide number of "
                      "physical devices ({})."
                  ).format(n1, device_count)
              )
          elif not n1_per_device:
              _device_count = ragged
              n1_per_device = 1
  
          return n1_per_device, _device_count
  
      def parallel_fn_x1(x1, x2=None, *args, **kwargs):
          x2_is_none = utils.all_none(x2)
          if x2_is_none:
              # TODO(schsam): Only compute the upper triangular part of the kernel.
              x2 = x1
  
          def get_batch_size(x):
              if utils.is_list_or_tuple(x):
                  return get_batch_size(x[0])
              return x.shape[0]
  
          n1 = get_batch_size(x1)
          n2 = n1 if x2_is_none else get_batch_size(x2)
  
          _check_dropout(n1, n2, kwargs)
          n1_per_device, _device_count = _get_n_per_device(n1)
  
          _kernel_fn = _jit_or_pmap_broadcast(kernel_fn, _device_count)
  
          @utils.nt_tree_fn()
          def batch_data(x):
              input_shape = x.shape[1:]
              return jnp.reshape(
                  x,
                  (
                      _device_count,
                      n1_per_device,
                  )
                  + input_shape,
              )
  
          for k, v in kwargs.items():
              if _is_np_ndarray(v):
                  assert isinstance(v, tuple) and len(v) == 2
                  v0 = jnp.reshape(
                      v[0],
                      (
                          _device_count,
                          n1_per_device,
                      )
                      + v[0].shape[1:],
                  )
                  kwargs[k] = (v0, v[1])
  
          x1 = batch_data(x1)
  
          kernel = _kernel_fn(x1, x2, *args, **kwargs)
          return _flatten_kernel(kernel, x2_is_none, True)
  
      def parallel_fn_kernel(kernel, *args, **kwargs):
          @utils.nt_tree_fn(reduce=lambda shapes: shapes[0])
          def get_batch_sizes(k):
              n1 = n2 = k.cov1.shape[0]
              if k.cov2 is not None:
                  n2 = k.cov2.shape[0]
              return n1, n2
  
          n1, n2 = get_batch_sizes(kernel)
          _check_dropout(n1, n2, kwargs)
          n1_per_device, _device_count = _get_n_per_device(n1)
  
          _kernel_fn = _jit_or_pmap_broadcast(kernel_fn, _device_count)
  
          cov2_is_none = utils.nt_tree_fn(reduce=lambda k: all(k))(
              lambda k: k.cov2 is None
          )(kernel)
          kernel = _reshape_kernel_for_pmap(
              kernel, _device_count, n1_per_device
          )
          kernel = _kernel_fn(kernel, *args, **kwargs)
          if cov2_is_none:
              kernel = _set_cov2_to_none(kernel)
          return _flatten_kernel(kernel, cov2_is_none, True)
  
      @utils.wraps(kernel_fn)
      def parallel_fn(x1_or_kernel, x2=None, *args, **kwargs):
          if utils.is_nt_tree_of(
              x1_or_kernel, (np.ndarray, jnp.ndarray)
          ):
              return parallel_fn_x1(x1_or_kernel, x2, *args, **kwargs)
          elif utils.is_nt_tree_of(x1_or_kernel, Kernel):
              assert not x2
              return parallel_fn_kernel(x1_or_kernel, *args, **kwargs)
          raise NotImplementedError()
  
      # Set function attributes so that `serial` can detect whether it is acting on
      # a parallel function.
      parallel_fn.device_count = device_count
      return parallel_fn
  
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 9 for calling a global function with example usage: _parallel(kernel_fn, use_serial, dropout_in_analytic_kernel, device_count) and returns: parallel_fn, n1_per_device,  _device_count, _flatten_kernel, _flatten_kernel, x, jnp, n1,  n2, parallel_fn_x1, get_batch_size, parallel_fn_kernel