code: def generate_batch_splits(
      samples_idx: np.ndarray, batch_size: int, drop_last=True
  ) -> np.ndarray:
      """Generate batches of data for a specified batch size from sample indices. If the dataset size is not divisible by
      the batch size and `drop_last` is `True`, the last incomplete batch is dropped. Else, it is returned.
      """
      num_samples = len(samples_idx)
      if drop_last:
          samples_to_remove = num_samples % batch_size
          if samples_to_remove != 0:
              samples_idx = samples_idx[:-samples_to_remove]
          sections_split = num_samples // batch_size
          samples_idx = samples_idx.reshape(
              (sections_split, batch_size)
          )
      else:
          sections_split = math.ceil(num_samples / batch_size)
          samples_idx = np.array_split(samples_idx, sections_split)
      return samples_idx
  
imports: import json
  import logging
  import math
  import os
  import sys
  import time
  import warnings
  from dataclasses import asdict, dataclass, field
  from enum import Enum
  from itertools import chain
  from pathlib import Path
  from typing import Dict, List, Optional, Tuple
  import flax
  import jax
  import jax.numpy
  import numpy
  import optax
  from datasets import load_dataset
  from flax import jax_utils, traverse_util
  from flax.jax_utils import pad_shard_unpad
  from flax.training import train_state
  from flax.training.common_utils import get_metrics, onehot, shard
  from huggingface_hub import Repository, create_repo
  from tqdm import tqdm
  from transformers import CONFIG_MAPPING, FLAX_MODEL_FOR_MASKED_LM_MAPPING, AutoConfig, AutoTokenizer, FlaxAutoModelForMaskedLM, HfArgumentParser, PreTrainedTokenizerBase, TensorType, is_tensorboard_available, set_seed
  from transformers.utils import send_example_telemetry
label_desc: Use image node_id 1 for calling a global function with example usage: generate_batch_splits(samples_idx, batch_size, drop_last) and returns: samples_idx