code: def graph_cross_validation(
      dataset: InMemoryDataset,
      n_splits: int = 10,
      shuffle: bool = True,
      random_seed: _typing.Optional[int] = ...,
      stratify: bool = False,
  ) -> InMemoryDataset:
      r"""Cross validation for graph classification data
  
      Parameters
      ----------
      dataset : InMemoryDataset
          dataset with multiple graphs.
  
      n_splits : int
          the number of folds to split.
  
      shuffle : bool
          shuffle or not for sklearn.model_selection.StratifiedKFold
  
      random_seed : int
          random_state for sklearn.model_selection.StratifiedKFold
  
      stratify: bool
      """
      if not isinstance(n_splits, int):
          raise TypeError
      elif not n_splits > 0:
          raise ValueError
      if not isinstance(shuffle, bool):
          raise TypeError
      if not (
          random_seed in (Ellipsis, None)
          or isinstance(random_seed, int)
      ):
          raise TypeError
      elif isinstance(random_seed, int) and random_seed >= 0:
          _random_seed: int = random_seed
      else:
          _random_seed: int = random.randrange(0, 65536)
      if not isinstance(stratify, bool):
          raise TypeError
  
      if stratify:
          kf = StratifiedKFold(
              n_splits=n_splits,
              shuffle=shuffle,
              random_state=_random_seed,
          )
      else:
          kf = KFold(
              n_splits=n_splits,
              shuffle=shuffle,
              random_state=_random_seed,
          )
      if _backend.DependentBackend.is_pyg():
          dataset_y = np.array([x["y"].item() for x in dataset])
      else:
          dataset_y = np.array([x[1].item() for x in dataset])
      idx_list = [
          (train_index.tolist(), test_index.tolist())
          for train_index, test_index in kf.split(
              np.zeros(len(dataset)), np.array(dataset_y)
          )
      ]
  
      # dataset.folds = idx_list
      setattr(dataset, "folds", idx_list)
      setattr(dataset, "train_index", idx_list[0][0])
      setattr(dataset, "val_index", idx_list[0][1])
      return dataset
  
file_path: AutoGL/autogl/datasets/utils/_general.py
imports: import numpy
  import random
  import torch
  import torch.utils.data
  import typing
  from sklearn.model_selection import StratifiedKFold, KFold
  from autogl import backend
  from autogl.data import InMemoryDataset
label_desc: Use image node_id 4 for calling a global function with example usage: graph_cross_validation(dataset, n_splits, shuffle, random_seed, stratify) and returns: dataset