code: def graph_get_split(
      dataset,
      mask: str = "train",
      is_loader: bool = True,
      batch_size: int = 128,
      num_workers: int = 0,
      shuffle: bool = False,
  ) -> _typing.Union[torch.utils.data.DataLoader, _typing.Iterable]:
      r"""Get train/test dataset/dataloader after cross validation.
  
      Parameters
      ----------
      dataset:
          dataset with multiple graphs.
  
      mask : str
  
      is_loader : bool
          return original dataset or data loader
  
      batch_size : int
          batch_size for generating Dataloader
      num_workers : int
          number of workers parameter for data loader
      shuffle: bool
          whether shuffle the dataloader
      """
      if not isinstance(mask, str):
          raise TypeError
      elif mask.lower() not in ("train", "val", "test"):
          raise ValueError
      if not isinstance(is_loader, bool):
          raise TypeError
      if not isinstance(batch_size, int):
          raise TypeError
      elif not batch_size > 0:
          raise ValueError
      if not isinstance(num_workers, int):
          raise TypeError
      elif not num_workers >= 0:
          raise ValueError
  
      if mask.lower() not in ("train", "val", "test"):
          raise ValueError
      elif mask.lower() == "train":
          try:
              optional_dataset_split = dataset.train_split
          except AttributeError:
              optional_dataset_split = dataset.cross_train_split
          if optional_dataset_split is None:
              raise ValueError(
                  f"Provided dataset do NOT have {mask} split"
              )
          else:
              sub_dataset = InMemoryDataset(
                  optional_dataset_split,
                  train_index=list(range(len(optional_dataset_split))),
              )
      elif mask.lower() == "val":
          try:
              optional_dataset_split = dataset.val_split
          except AttributeError:
              optional_dataset_split = dataset.cross_val_split
          if optional_dataset_split is None:
              raise ValueError(
                  f"Provided dataset do NOT have {mask} split"
              )
          else:
              sub_dataset = InMemoryDataset(
                  optional_dataset_split,
                  val_index=list(range(len(optional_dataset_split))),
              )
      elif mask.lower() == "test":
          optional_dataset_split = dataset.test_split
          if optional_dataset_split is None:
              raise ValueError(
                  f"Provided dataset do NOT have {mask} split"
              )
          else:
              sub_dataset = InMemoryDataset(
                  optional_dataset_split,
                  test_index=list(range(len(optional_dataset_split))),
              )
      else:
          raise ValueError(
              f"The provided mask parameter must be a str in ['train', 'val', 'test'], "
              f"illegal provided value is [{mask}]"
          )
      if not is_loader:
          return sub_dataset
      if is_loader:
          if not (
              _backend.DependentBackend.is_dgl()
              or _backend.DependentBackend.is_pyg()
          ):
              raise RuntimeError("Unsupported backend")
          elif _backend.DependentBackend.is_dgl():
              from dgl.dataloading import GraphDataLoader
  
              return GraphDataLoader(
                  sub_dataset,
                  **{
                      "batch_size": batch_size,
                      "num_workers": num_workers,
                  },
                  shuffle=shuffle,
              )
          elif _backend.DependentBackend.is_pyg():
              _sub_dataset: _typing.Any = optional_dataset_split
              import torch_geometric
  
              if int(torch_geometric.__version__.split(".")[0]) >= 2:
                  # version 2.x
                  from torch_geometric.loader import DataLoader
              else:
                  from torch_geometric.data import DataLoader
              return DataLoader(
                  _sub_dataset,
                  batch_size=batch_size,
                  num_workers=num_workers,
                  shuffle=shuffle,
              )
      else:
          return sub_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 7 for calling a global function with example usage: graph_get_split(dataset, mask, is_loader, batch_size, num_workers, shuffle) and returns: sub_dataset, sub_dataset, GraphDataLoader, DataLoader