code: def _serial(
      kernel_fn: _KernelFn,
      batch_size: int,
      store_on_device: bool = True,
  ) -> _KernelFn:
      """Returns a function that computes a kernel in batches serially.
  
      This function computes the kernel over data in batches where each batch is
      processed sequentially with a given batch size. If serial detects that the
      kernel function is the result of `_parallel` (that is, if the kernel is
      distributed over multiple devices) then serial adjusts the batch size so that
      each device processes chunks of work that have batch_size x batch_size.
  
      The dataset size must divide the effective batch size. If parallelism is used
      this means that `|x1|` must divide `batch_size * device_count` and `|x2|` must
      divide `batch_size`.
  
      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`.
  
        batch_size:
          Integer specifying the size of batches in which to split the data.
  
        store_on_device:
          A boolean that species whether the computed kernel should be kept on
          device or brought back to CPU as it is computed. Defaults to `True`.
  
      Returns:
        A new function with the same signature as kernel_fn that computes the kernel
        by batching over the dataset serially with the specified batch_size.
      """
  
      device_count = max(getattr(kernel_fn, "device_count", 1), 1)
  
      if not store_on_device:
          _kernel_fn = kernel_fn
  
          @utils.wraps(_kernel_fn)
          def kernel_fn(x1, x2=None, *args, **kwargs):
              return device_put(
                  _kernel_fn(x1, x2, *args, **kwargs), devices("cpu")[0]
              )
  
      flatten = partial(_flatten_kernel, is_parallel=False)
  
      def serial_fn_x1(
          x1: NTTree[jnp.ndarray],
          x2: Optional[NTTree[Optional[jnp.ndarray]]] = None,
          *args,
          **kwargs,
      ) -> NTTree[Kernel]:
          x2_is_none = utils.all_none(x2)
          if x2_is_none:
              # TODO(schsam): Only compute the upper triangular part of the kernel.
              x2 = x1
  
          @utils.nt_tree_fn(reduce=lambda x: x[0])
          def get_n1_n2(x1, x2):
              n1, n2 = x1.shape[0], x2.shape[0]
              return n1, n2
  
          n1, n2 = get_n1_n2(x1, x2)
  
          (
              n1_batches,
              n1_batch_size,
              n2_batches,
              n2_batch_size,
          ) = _get_n_batches_and_batch_sizes(
              n1, n2, batch_size, device_count
          )
  
          @utils.nt_tree_fn(nargs=1)
          def batch_input(
              x: jnp.ndarray, batch_count: int, batch_size: int
          ):
              input_shape = x.shape[1:]
              return jnp.reshape(
                  x,
                  (
                      batch_count,
                      batch_size,
                  )
                  + input_shape,
              )
  
          x1s = batch_input(x1, n1_batches, n1_batch_size)
          x2s = batch_input(x2, n2_batches, n2_batch_size)
  
          kwargs_np1 = {}
          kwargs_np2 = {}
          kwargs_other = {}
  
          for k, v in kwargs.items():
              if _is_np_ndarray(v):
                  if k == "rng":
                      key1, key2 = random.split(v)
                      v1 = random.split(key1, n1_batches)
                      v2 = random.split(key2, n2_batches)
                  else:
                      assert isinstance(v, tuple) and len(v) == 2
                      v1 = jnp.reshape(
                          v[0],
                          (
                              n1_batches,
                              n1_batch_size,
                          )
                          + v[0].shape[1:],
                      )
                      v2 = jnp.reshape(
                          v[1],
                          (
                              n2_batches,
                              n2_batch_size,
                          )
                          + v[1].shape[1:],
                      )
                  kwargs_np1[k] = v1
                  kwargs_np2[k] = v2
              else:
                  kwargs_other[k] = v
  
          def row_fn(_, x1):
              return _, _scan(col_fn, x1, (x2s, kwargs_np2))[1]
  
          def col_fn(x1, x2):
              x1, kwargs1 = x1
              x2, kwargs2 = x2
              kwargs_merge = {
                  **kwargs_other,
                  **dict(
                      (k, (kwargs1[k], kwargs2[k])) for k in kwargs1
                  ),
              }
              return (x1, kwargs1), kernel_fn(
                  x1, x2, *args, **kwargs_merge
              )
  
          _, kernel = _scan(row_fn, 0, (x1s, kwargs_np1))
          return flatten(kernel, x2_is_none)
  
      def serial_fn_kernel(
          k: NTTree[Kernel], *args, **kwargs
      ) -> NTTree[Kernel]:
          def get_n1_n2(k: NTTree[Kernel]) -> tuple[int, ...]:
              if utils.is_list_or_tuple(k):
                  # TODO(schsam): We might want to check for consistency here, but I can't
                  # imagine a case where we could get inconsistent kernels.
                  return get_n1_n2(k[0])
  
              if isinstance(k, Kernel):
                  return k.nngp.shape[
                      :2
                  ]  # pytype: disable=attribute-error
  
              raise TypeError(type(Kernel), Kernel)
  
          n1, n2 = get_n1_n2(k)
  
          (
              n1_batches,
              n1_batch_size,
              n2_batches,
              n2_batch_size,
          ) = _get_n_batches_and_batch_sizes(
              n1, n2, batch_size, device_count
          )
  
          n1s = jnp.arange(0, n1, n1_batch_size)
          n2s = jnp.arange(0, n2, n2_batch_size)
  
          @utils.nt_tree_fn(nargs=1)
          def slice_kernel(k, n1, n2):
              return k.slice(n1, n2)
  
          kwargs_np1 = {}
          kwargs_np2 = {}
  
          kwargs_other = {}
          for key, v in kwargs.items():
              if _is_np_ndarray(v):
                  assert isinstance(v, tuple) and len(v) == 2
                  v1 = jnp.reshape(
                      v[0],
                      (
                          n1_batches,
                          n1_batch_size,
                      )
                      + v[0].shape[1:],
                  )
                  v2 = jnp.reshape(
                      v[1],
                      (
                          n2_batches,
                          n2_batch_size,
                      )
                      + v[1].shape[1:],
                  )
                  kwargs_np1[key] = v1
                  kwargs_np2[key] = v2
              else:
                  kwargs_other[key] = v
  
          def row_fn(_, n1):
              return _, _scan(col_fn, n1, (n2s, kwargs_np2))[1]
  
          def col_fn(n1, n2):
              # NOTE(schsam): If we end up wanting to enable jit-of-batch then we will
              # probably have to change this to dynamic slicing.
              n1, kwargs1 = n1
              n2, kwargs2 = n2
              kwargs_merge = {
                  **kwargs_other,
                  **dict(
                      (key, (kwargs1[key], kwargs2[key]))
                      for key in kwargs1
                  ),
              }
              n1_slice = slice(n1, n1 + n1_batch_size)
              n2_slice = slice(n2, n2 + n2_batch_size)
              in_kernel = slice_kernel(k, n1_slice, n2_slice)
              return (n1, kwargs1), kernel_fn(
                  in_kernel, *args, **kwargs_merge
              )
  
          cov2_is_none = utils.nt_tree_fn(reduce=all)(
              lambda k: k.cov2 is None
          )(k)
          _, k = _scan(row_fn, 0, (n1s, kwargs_np1))
          if cov2_is_none:
              k = _set_cov2_to_none(k)
          return flatten(k, cov2_is_none)
  
      @utils.wraps(kernel_fn)
      def serial_fn(
          x1_or_kernel: Union[NTTree[jnp.ndarray], NTTree[Kernel]],
          x2: Optional[NTTree[Optional[jnp.ndarray]]] = None,
          *args,
          **kwargs,
      ) -> NTTree[Kernel]:
          if utils.is_nt_tree_of(
              x1_or_kernel, (np.ndarray, jnp.ndarray)
          ):
              return serial_fn_x1(x1_or_kernel, x2, *args, **kwargs)
          elif utils.is_nt_tree_of(x1_or_kernel, Kernel):
              if x2 is not None:
                  raise ValueError(f"`x2` must be `None`, got {x2}.")
              return serial_fn_kernel(x1_or_kernel, *args, **kwargs)
          else:
              raise TypeError(x1_or_kernel, type(x1_or_kernel))
  
      return serial_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 8 for calling a global function with example usage: _serial(kernel_fn, batch_size, store_on_device) and returns: serial_fn, flatten, flatten, device_put, n1,  n2, jnp, _,  _scan, x1,  kwargs1,  kernel_fn, k, _,  _scan, n1,  kwargs1,  kernel_fn, serial_fn_x1, get_n1_n2, k, serial_fn_kernel