class_docstr: Partition the graph structure and its features of a
  :class:`~torch_geometric.data.Data` or
  :class:`~torch_geometric.data.HeteroData` object.
  
  Partitioned data output will be structured like this:
  
  **Homogeneous graphs:**
  
  .. code-block::
  
      root/
      |-- META.json
      |-- node_map.pt
      |-- edge_map.pt
      |-- part0/
          |-- graph.pt
          |-- node_feats.pt
          |-- edge_feats.pt
      |-- part1/
          |-- graph.pt
          |-- node_feats.pt
          |-- edge_feats.pt
  
  **Heterogeneous graphs:**
  
  .. code-block::
  
      root/
      |-- META.json
      |-- node_map/
          |-- ntype1.pt
          |-- ntype2.pt
      |-- edge_map/
          |-- etype1.pt
          |-- etype2.pt
      |-- part0/
          |-- graph.pt
          |-- node_feats.pt
          |-- edge_feats.pt
      |-- part1/
          |-- graph.pt
          |-- node_feats.pt
          |-- edge_feats.pt
  
  Args:
      data (Data or HeteroData): The data object.
      num_parts (int): The number of partitions.
      recursive (bool, optional): If set to :obj:`True`, will use multilevel
          recursive bisection instead of multilevel k-way partitioning.
          (default: :obj:`False`)
      root (str): Root directory where the partitioned dataset should be
          saved.
class_docstr_tok: ["Partition","the","graph","structure","and","its","features","of","a",":","class",":","`","~torch_geometric.data.Data","`","or",":","class",":","`","~torch_geometric.data.HeteroData","`","object",".","Partitioned","data","output","will","be","structured","like","this",":","*","*","Homogeneous","graphs",":","*","*","..","code-block",":",":","root\/","|","--","META.json","|","--","node_map.pt","|","--","edge_map.pt","|","--","part0\/","|","--","graph.pt","|","--","node_feats.pt","|","--","edge_feats.pt","|","--","part1\/","|","--","graph.pt","|","--","node_feats.pt","|","--","edge_feats.pt","*","*","Heterogeneous","graphs",":","*","*","..","code-block",":",":","root\/","|","--","META.json","|","--","node_map\/","|","--","ntype1.pt","|","--","ntype2.pt","|","--","edge_map\/","|","--","etype1.pt","|","--","etype2.pt","|","--","part0\/","|","--","graph.pt","|","--","node_feats.pt","|","--","edge_feats.pt","|","--","part1\/","|","--","graph.pt","|","--","node_feats.pt","|","--","edge_feats.pt","Args",":","data","(","Data","or","HeteroData",")",":","The","data","object",".","num_parts","(","int",")",":","The","number","of","partitions",".","recursive","(","bool",",","optional",")",":","If","set","to",":","obj",":","`","True","`",",","will","use","multilevel","recursive","bisection","instead","of","multilevel","k-way","partitioning",".","(","default",":",":","obj",":","`","False","`",")","root","(","str",")",":","Root","directory","where","the","partitioned","dataset","should","be","saved","."]
docstr: Generates the partition.
docstr_tok: ["Generates","the","partition","."]
code: def generate_partition(self):
      r"""Generates the partition."""
      os.makedirs(self.root, exist_ok=True)
      data = self.data.to_homogeneous() if self.is_hetero else self.data
      cluster_data = ClusterData(
          data,
          num_parts=self.num_parts,
          recursive=self.recursive,
          log=True,
          keep_inter_cluster_edges=True,
          sparse_format="csc",
      )
  
      node_perm = cluster_data.partition.node_perm
      partptr = cluster_data.partition.partptr
      edge_perm = cluster_data.partition.edge_perm
  
      node_map = torch.empty(data.num_nodes, dtype=torch.int64)
      edge_map = torch.empty(data.num_edges, dtype=torch.int64)
      node_offset, edge_offset = {}, {}
  
      if self.is_hetero:
          offset = 0
          for node_type in self.node_types:
              node_offset[node_type] = offset
              offset += self.data[node_type].num_nodes
  
          offset = 0
          for edge_name in self.edge_types:
              edge_offset[edge_name] = offset
              offset += self.data.num_edges_dict[edge_name]
  
          edge_start = 0
          for pid in range(self.num_parts):
              logging.info(f"Saving graph partition {pid}")
              path = osp.join(self.root, f"part_{pid}")
              os.makedirs(path, exist_ok=True)
  
              part_data = cluster_data[pid]
              start, end = int(partptr[pid]), int(partptr[pid + 1])
  
              num_edges = part_data.num_edges
              edge_id = edge_perm[edge_start : edge_start + num_edges]
              edge_map[edge_id] = pid
              edge_start += num_edges
  
              node_id = node_perm[start:end]
              node_map[node_id] = pid
  
              graph = {}
              efeat = {}
              for i, edge_type in enumerate(self.edge_types):
                  # Row vector refers to source nodes.
                  # Column vector refers to destination nodes.
                  src, _, dst = edge_type
                  size = (
                      self.data[src].num_nodes,
                      self.data[dst].num_nodes,
                  )
  
                  mask = part_data.edge_type == i
                  row = part_data.edge_index[0, mask]
                  col = part_data.edge_index[1, mask]
                  global_col = node_id[col]
                  global_row = node_perm[row]
  
                  # Sort on col to avoid keeping track of permuations in
                  # NeighborSampler when converting to CSC format:
                  num_cols = col.size()[0]
                  global_col, perm = index_sort(
                      global_col, max_value=num_cols
                  )
                  global_row = global_row[perm]
                  global_eid = edge_id[mask][perm]
                  assert torch.equal(
                      data.edge_index[:, global_eid],
                      torch.stack((global_row, global_col), dim=0),
                  )
                  local_eid = global_eid - edge_offset[edge_type]
                  assert torch.equal(
                      self.data[edge_type].edge_index[:, local_eid],
                      torch.stack(
                          (
                              global_row - node_offset[src],
                              global_col - node_offset[dst],
                          ),
                          dim=0,
                      ),
                  )
                  graph[edge_type] = {
                      "edge_id": global_eid,
                      "row": global_row - node_offset[src],
                      "col": global_col - node_offset[dst],
                      "size": size,
                  }
  
                  if "edge_attr" in part_data:
                      edge_attr = part_data.edge_attr[mask][perm]
                      efeat[edge_type] = {
                          "global_id": global_eid,
                          "feats": dict(edge_attr=edge_attr),
                      }
  
              torch.save(efeat, osp.join(path, "edge_feats.pt"))
              torch.save(graph, osp.join(path, "graph.pt"))
  
              nfeat = {}
              for i, node_type in enumerate(self.node_types):
                  mask = part_data.node_type == i
                  x = part_data.x[mask] if "x" in part_data else None
                  nfeat[node_type] = {
                      "global_id": node_id[mask],
                      "id": node_id[mask] - node_offset[node_type],
                      "feats": dict(x=x),
                  }
              torch.save(nfeat, osp.join(path, "node_feats.pt"))
  
          logging.info("Saving partition mapping info")
          path = osp.join(self.root, "node_map")
          os.makedirs(path, exist_ok=True)
          for i, node_type in enumerate(self.node_types):
              mask = data.node_type == i
              torch.save(
                  node_map[mask], osp.join(path, f"{node_type}.pt")
              )
  
          path = osp.join(self.root, "edge_map")
          os.makedirs(path, exist_ok=True)
          for i, edge_type in enumerate(self.edge_types):
              mask = data.edge_type == i
              torch.save(
                  edge_map[mask],
                  osp.join(path, f"{EdgeTypeStr(edge_type)}.pt"),
              )
  
      else:  # `if not self.is_hetero:`
          edge_start = 0
          for pid in range(self.num_parts):
              logging.info(f"Saving graph partition {pid}")
              path = osp.join(self.root, f"part_{pid}")
              os.makedirs(path, exist_ok=True)
  
              part_data = cluster_data[pid]
              start, end = int(partptr[pid]), int(partptr[pid + 1])
  
              num_edges = part_data.num_edges
              edge_id = edge_perm[edge_start : edge_start + num_edges]
              edge_map[edge_id] = pid
              edge_start += num_edges
  
              node_id = node_perm[start:end]  # global node_ids
              node_map[node_id] = pid  # 0 or 1
  
              row = part_data.edge_index[0]
              col = part_data.edge_index[1]
              num_cols = col.size()[0]
  
              global_col = node_id[col]  # part_ids -> global
              global_row = node_perm[row]
  
              # Sort on col to avoid keeping track of permuations in
              # NeighborSampler when converting to CSC format:
              global_col, perm = index_sort(
                  global_col, max_value=num_cols
              )
              global_row = global_row[perm]
              edge_id = edge_id[perm]
  
              assert torch.equal(
                  self.data.edge_index[:, edge_id],
                  torch.stack((global_row, global_col)),
              )
              if "edge_attr" in part_data:
                  edge_attr = part_data.edge_attr[perm]
                  assert torch.equal(
                      self.data.edge_attr[edge_id, :], edge_attr
                  )
  
              torch.save(
                  {
                      "edge_id": edge_id,
                      "row": global_row,
                      "col": global_col,
                      "size": (data.num_nodes, data.num_nodes),
                  },
                  osp.join(path, "graph.pt"),
              )
  
              torch.save(
                  {
                      "global_id": node_id,
                      "feats": dict(x=part_data.x),
                  },
                  osp.join(path, "node_feats.pt"),
              )
              if "edge_attr" in part_data:
                  torch.save(
                      {
                          "global_id": edge_id,
                          "feats": dict(
                              edge_attr=part_data.edge_attr[perm]
                          ),
                      },
                      osp.join(path, "edge_feats.pt"),
                  )
  
          logging.info("Saving partition mapping info")
          torch.save(node_map, osp.join(self.root, "node_map.pt"))
          torch.save(edge_map, osp.join(self.root, "edge_map.pt"))
  
      logging.info("Saving metadata")
      meta = {
          "num_parts": self.num_parts,
          "node_types": self.node_types,
          "edge_types": self.edge_types,
          "node_offset": list(node_offset.values())
          if node_offset
          else None,
          "is_hetero": self.is_hetero,
          "is_sorted": True,  # Based on colum/destination.
      }
      with open(osp.join(self.root, "META.json"), "w") as f:
          json.dump(meta, f)
  
file_path: pytorch_geometric/torch_geometric/distributed/partition.py
imports: import json
  import logging
  import os
  import os.path
  from typing import List, Optional, Union
  import torch
  from torch_geometric.data import Data, HeteroData
  from torch_geometric.distributed.utils import as_str
  from torch_geometric.loader.cluster import ClusterData
  from torch_geometric.typing import Dict, EdgeType, EdgeTypeStr, NodeType, Tuple
  from torch_geometric.utils import index_sort
label_desc: Use image node_id 5 for calling the Partitioner obj's underlying member method code with example usage: obj.generate_partition() without return types