code: def generate_batch_splits(
      samples_idx: np.ndarray, batch_size: int
  ) -> np.ndarray:
      nb_samples = len(samples_idx)
      samples_to_remove = nb_samples % batch_size
  
      if samples_to_remove != 0:
          samples_idx = samples_idx[:-samples_to_remove]
      sections_split = nb_samples // batch_size
      batch_idx = np.split(samples_idx, sections_split)
      return batch_idx
  
imports: import logging
  import os
  import sys
  from dataclasses import dataclass, field
  from pathlib import Path
  from typing import Dict, List, Optional, Tuple
  import numpy
  from datasets import load_dataset
  from tqdm import tqdm
  import jax
  import jax.numpy
  from flax import jax_utils
  from flax.optim import Adam
  from flax.training import common_utils
  from flax.training.common_utils import get_metrics
  from jax.nn import log_softmax
  from modeling_flax_performer import FlaxPerformerForMaskedLM
  from transformers import MODEL_FOR_MASKED_LM_MAPPING, AutoTokenizer, BertConfig, FlaxBertForMaskedLM, HfArgumentParser, PreTrainedTokenizerBase, TensorType, TrainingArguments, is_tensorboard_available, set_seed
label_desc: Use image node_id 7 for calling a global function with example usage: generate_batch_splits(samples_idx, batch_size) and returns: batch_idx