code: def _gen_train_test_split(
      folds: int, arm_names: np.ndarray
  ) -> Iterable[Tuple[Set[str], Set[str]]]:
      """Return train/test splits of arm names.
  
      Args:
          folds: Number of folds to return
          arm_names: Array of arm names
  
      Returns:
          Yields (train, test) tuple of arm names.
      """
      n = len(arm_names)
      test_size = n // folds  # The size of all test sets but the last
      final_size = test_size + (
          n - folds * test_size
      )  # Grab the leftovers
      for fold in range(folds):
          # We will take the test set from the back of the array.
          # Roll the list of arm names to get a fresh test set
          arm_names = np.roll(arm_names, test_size)
          n_test = test_size if fold < folds - 1 else final_size
          yield set(arm_names[:-n_test]), set(arm_names[-n_test:])
  
imports: from abc import ABC, abstractmethod
  from collections import defaultdict
  from enum import Enum
  from functools import partial
  from logging import Logger
  from numbers import Number
  from typing import Any, Callable, Dict, Iterable, List, NamedTuple, Optional, Set, Tuple
  import numpy
  from ax.core.observation import Observation, ObservationData
  from ax.core.optimization_config import OptimizationConfig
  from ax.modelbridge.base import ModelBridge
  from ax.utils.common.logger import get_logger
  from ax.utils.stats.model_fit_stats import _correlation_coefficient, _fisher_exact_test_p, _log_likelihood, _mape, _mean_prediction_ci, _rank_correlation, _total_raw_effect, _wmape, compute_model_fit_metrics
label_desc: Use image node_id 7 for calling a global function with example usage: _gen_train_test_split(folds, arm_names) without return types