code: def load_partition_info(
      root_dir: str,
      partition_idx: int,
  ) -> Tuple[Dict, int, int, torch.Tensor, torch.Tensor]:
      # load the partition with PyG format (graphstore/featurestore)
      with open(osp.join(root_dir, "META.json"), "rb") as infile:
          meta = json.load(infile)
      num_partitions = meta["num_parts"]
      assert partition_idx >= 0
      assert partition_idx < num_partitions
      partition_dir = osp.join(root_dir, f"part_{partition_idx}")
      assert osp.exists(partition_dir)
  
      if meta["is_hetero"] is False:
          node_pb = torch.load(osp.join(root_dir, "node_map.pt"))
          edge_pb = torch.load(osp.join(root_dir, "edge_map.pt"))
  
          return (meta, num_partitions, partition_idx, node_pb, edge_pb)
      else:
          node_pb_dict = {}
          node_pb_dir = osp.join(root_dir, "node_map")
          for ntype in meta["node_types"]:
              node_pb_dict[ntype] = torch.load(
                  osp.join(node_pb_dir, f"{as_str(ntype)}.pt")
              )
  
          edge_pb_dict = {}
          edge_pb_dir = osp.join(root_dir, "edge_map")
          for etype in meta["edge_types"]:
              edge_pb_dict[tuple(etype)] = torch.load(
                  osp.join(edge_pb_dir, f"{as_str(etype)}.pt")
              )
  
          return (
              meta,
              num_partitions,
              partition_idx,
              node_pb_dict,
              edge_pb_dict,
          )
  
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 1 for calling a global function with example usage: load_partition_info(root_dir, partition_idx) and returns: meta,  num_partitions,  partition_idx,  node_pb,  edge_pb, meta,  num_partitions,  partition_idx,  node_pb_dict,  edge_pb_dict