code: def generate_batch_splits(
      samples_idx: np.ndarray, batch_size: int
  ) -> np.ndarray:
      num_samples = len(samples_idx)
      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
      batch_idx = np.split(samples_idx, sections_split)
      return batch_idx
  
imports: import logging
  import sys
  import time
  from dataclasses import field
  from pathlib import Path
  from typing import Dict, List, Optional, Union
  import numpy
  from datasets import DatasetDict, load_dataset
  from tqdm import tqdm
  import flax
  import jax
  import jax.numpy
  import librosa
  import optax
  from flax import jax_utils, traverse_util
  from flax.training import train_state
  from flax.training.common_utils import get_metrics, onehot, shard
  from transformers import FlaxWav2Vec2ForPreTraining, HfArgumentParser, TrainingArguments, Wav2Vec2Config, Wav2Vec2FeatureExtractor, is_tensorboard_available
  from transformers.models.wav2vec2.modeling_flax_wav2vec2 import _compute_mask_indices, _sample_negative_indices
label_desc: Use image node_id 4 for calling a global function with example usage: generate_batch_splits(samples_idx, batch_size) and returns: batch_idx