code: def random_splits_mask_class(
      dataset: InMemoryDataset,
      num_train_per_class: int = 20,
      num_val_per_class: int = 30,
      seed: _typing.Optional[int] = ...,
  ):
      r"""If the data has masks for train/val/test, return the splits with specific number of samples from every class for training as suggested in Pitfalls of graph neural network evaluation [#]_ for semi-supervised learning.
  
      References
      ----------
      .. [#] Shchur, O., Mumme, M., Bojchevski, A., & Günnemann, S. (2018).
          Pitfalls of graph neural network evaluation.
          arXiv preprint arXiv:1811.05868.
  
      Parameters
      ----------
      dataset: InMemoryDataset
          instance of ``InMemoryDataset``
      num_train_per_class : int
          the number of samples from every class used for training.
      num_val_per_class : int
          the number of samples from every class used for validation.
      seed : int
          random seed for splitting dataset.
      """
  
      if _backend.DependentBackend.is_pyg():
          import torch_geometric.transforms as T
  
          transform = T.RandomNodeSplit(
              split="test_rest",
              num_train_per_class=num_train_per_class,
              num_val=num_val_per_class,
          )
          dataset = [transform(data) for data in dataset]
      else:
          import dgl
  
          def __split_data(data):
              if "label" not in data.ndata:
                  return data
              elif "label" in data.ndata:
                  label: torch.Tensor = data.ndata["label"]
              else:
                  raise RuntimeError
              num_nodes: int = label.size(0)
              num_classes: int = label.cpu().max().item() + 1
  
              _rng_state: torch.Tensor = torch.get_rng_state()
              if seed not in (Ellipsis, None) and isinstance(seed, int):
                  torch.manual_seed(seed)
              train_mask = torch.zeros(
                  num_nodes, dtype=torch.bool, device=label.device
              )
              val_mask = torch.zeros(
                  num_nodes, dtype=torch.bool, device=label.device
              )
              test_mask = torch.zeros(
                  num_nodes, dtype=torch.bool, device=label.device
              )
              for class_index in range(num_classes):
                  idx = (label == class_index).nonzero().view(-1)
                  assert (
                      num_train_per_class + num_val_per_class
                      < idx.size(0)
                  ), (
                      f"the total number of samples from every class "
                      f"used for training and validation is larger than "
                      f"the total samples in class [{class_index}]"
                  )
                  randomized_index: torch.Tensor = torch.randperm(
                      idx.size(0)
                  )
                  train_idx = idx[
                      randomized_index[:num_train_per_class]
                  ]
                  val_idx = idx[
                      randomized_index[
                          num_train_per_class : (
                              num_train_per_class + num_val_per_class
                          )
                      ]
                  ]
                  train_mask[train_idx] = True
                  val_mask[val_idx] = True
              else:
                  remaining = (
                      (~(train_mask + val_mask)).nonzero().view(-1)
                  )
                  test_mask[remaining] = True
  
              torch.set_rng_state(_rng_state)
              data.ndata["train_mask"] = train_mask
              data.ndata["val_mask"] = val_mask
              data.ndata["test_mask"] = test_mask
              if not torch.all(data.in_degrees() == 0):
                  data = dgl.add_self_loop(data)
              return data
  
          new_dataset = []
          for graph_index in range(len(dataset)):
              new_dataset.append(__split_data(dataset[graph_index]))
          dataset = new_dataset
      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 3 for calling a global function with example usage: random_splits_mask_class(dataset, num_train_per_class, num_val_per_class, seed) and returns: dataset, data, data