code: def batch(
      X: np.ndarray,
      batch_size: Optional[int] = None,
      n_batches: int = 4,
  ) -> List[np.ndarray]:
      """
      Splits the input into sub-arrays.
  
      Parameters
      ----------
      X
          Array to be split.
      batch_size
          The size of each batch. In particular
  
           - if `batch_size` is not ``None``, batches of this size are created. The sizes of the batches created might \
          vary if the 0-th dimension of `X` is not divisible by `batch_size`. For an array of length `l` that should \
          be split into `n` sections, it returns `l % n` sub-arrays of size `l//n + 1` and the rest of  `size l//n`
          
           - if `batch_size` is ``None``, then `X` is split into `n_batches` sub-arrays.
              
      n_batches
          Number of batches in which to split the sub-array. Only used if ``batch_size = None``
  
      Returns
      ------
      A list of sub-arrays of `X`.
      """  # noqa W605
  
      n_records = X.shape[0]
      if isinstance(X, sparse.spmatrix):
          logger.warning(
              "Batching function received sparse matrix input. Converting to dense matrix first..."
          )
          X = X.toarray()
  
      if batch_size:
          n_batches = n_records // batch_size
          if n_records % batch_size != 0:
              n_batches += 1
          slices = [batch_size * i for i in range(1, n_batches)]
          batches = np.array_split(X, slices)
      else:
          batches = np.array_split(X, n_batches)
  
      return batches
  
imports: import copy
  import logging
  from functools import partial
  from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union
  import numpy
  from scipy import sparse
  import ray
label_desc: Use image node_id 1 for calling a global function with example usage: batch(X, batch_size, n_batches) and returns: batches