def batch_array(array, max_batch_size):
  """Split an array according the maximum batch_size.

  Args:
    array: List or 2D Numpy array.
    max_batch_size: Integer value with maximum batch_size or None.

  Returns:
    A list of lists or numpy arrays the concatenation of which is the input
    list or array, and the first dimension of each array is less than or equal
    to max_batch_size. If max_batch_size is None, then a singleton list with the
    input list/array is returned.
  """
  if max_batch_size is None:
    return [array]
  num_batches = (len(array) + max_batch_size - 1) // max_batch_size
  assert num_batches > 0
  return [
      array[(i * max_batch_size):((i + 1) * max_batch_size)]
      for i in range(num_batches)
  ]