code: def _get_n_batches_and_batch_sizes(
      n1: int, n2: int, batch_size: int, device_count: int
  ) -> tuple[int, int, int, int]:
      # TODO(romann): if dropout batching works for different batch sizes, relax.
      max_serial_batch_size = np.gcd(n1, n2) // device_count
  
      n2_batch_size = min(batch_size, max_serial_batch_size)
      if n2_batch_size != batch_size:
          warnings.warn(
              "Batch size is reduced from requested %d to effective %d to "
              "fit the dataset." % (batch_size, n2_batch_size)
          )
  
      n1_batch_size = n2_batch_size * device_count
      n1_batches, ragged = divmod(n1, n1_batch_size)
      if ragged:
          # TODO(schsam): Relax this constraint.
          msg = (
              "Number of rows of kernel must divide batch size. Found n1 = {} "
              "and batch size = {}."
          ).format(n1, n1_batch_size)
          if device_count > 1:
              msg += (
                  " Note that device parallelism was detected and so the batch "
                  "size was expanded by a factor of {}.".format(
                      device_count
                  )
              )
          raise ValueError(msg)
  
      n2_batches, ragged = divmod(n2, n2_batch_size)
      if ragged:
          # TODO(schsam): Relax this constraint.
          raise ValueError(
              (
                  "Number of columns of kernel must divide batch "
                  "size. Found n2 = {} "
                  "and batch size = {}"
              ).format(n2, n2_batch_size)
          )
      return n1_batches, n1_batch_size, n2_batches, n2_batch_size
  
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 10 for calling a global function with example usage: _get_n_batches_and_batch_sizes(n1, n2, batch_size, device_count) and returns: n1_batches,  n1_batch_size,  n2_batches,  n2_batch_size