code: def partition_dataset(
      dataset_name: str,
      root_dir: str,
      num_parts: int,
      recursive: bool = False,
      use_sparse_tensor: bool = False,
  ):
      if not osp.isabs(root_dir):
          path = osp.dirname(osp.realpath(__file__))
          root_dir = osp.join(path, root_dir)
  
      dataset_dir = osp.join(root_dir, "dataset", dataset_name)
      dataset = get_dataset(
          dataset_name, dataset_dir, use_sparse_tensor
      )
      data = dataset[0]
  
      save_dir = osp.join(
          root_dir, "partitions", dataset_name, f"{num_parts}-parts"
      )
  
      partitions_dir = osp.join(save_dir, f"{dataset_name}-partitions")
      partitioner = Partitioner(
          data, num_parts, partitions_dir, recursive
      )
      partitioner.generate_partition()
  
      print("-- Saving label ...")
      label_dir = osp.join(save_dir, f"{dataset_name}-label")
      os.makedirs(label_dir, exist_ok=True)
  
      split_data = data["paper"] if dataset_name == "ogbn-mag" else data
  
      torch.save(
          split_data.y.squeeze(), osp.join(label_dir, "label.pt")
      )
  
      train_idx, val_idx, test_idx = get_idx_split(
          dataset, dataset_name, split_data
      )
  
      print("-- Partitioning training indices ...")
      train_idx = train_idx.split(train_idx.size(0) // num_parts)
      train_part_dir = osp.join(
          save_dir, f"{dataset_name}-train-partitions"
      )
      os.makedirs(train_part_dir, exist_ok=True)
      for i in range(num_parts):
          torch.save(
              train_idx[i], osp.join(train_part_dir, f"partition{i}.pt")
          )
  
      print("-- Partitioning validation indices ...")
      val_idx = val_idx.split(val_idx.size(0) // num_parts)
      val_part_dir = osp.join(
          save_dir, f"{dataset_name}-val-partitions"
      )
      os.makedirs(val_part_dir, exist_ok=True)
      for i in range(num_parts):
          torch.save(
              val_idx[i], osp.join(val_part_dir, f"partition{i}.pt")
          )
  
      print("-- Partitioning test indices ...")
      test_idx = test_idx.split(test_idx.size(0) // num_parts)
      test_part_dir = osp.join(
          save_dir, f"{dataset_name}-test-partitions"
      )
      os.makedirs(test_part_dir, exist_ok=True)
      for i in range(num_parts):
          torch.save(
              test_idx[i], osp.join(test_part_dir, f"partition{i}.pt")
          )
  
file_path: pytorch_geometric/examples/distributed/pyg/partition_graph.py
imports: import argparse
  import os
  import os.path
  import torch
  from ogb.nodeproppred import PygNodePropPredDataset
  import torch_geometric.transforms
  from torch_geometric.datasets import OGB_MAG, Reddit
  from torch_geometric.distributed import Partitioner
  from torch_geometric.utils import mask_to_index
label_desc: Use image node_id 1 for calling a global function with example usage: partition_dataset(dataset_name, root_dir, num_parts, recursive, use_sparse_tensor) without return types