code: def test_partition_hetero_data(tmp_path):
      data = FakeHeteroDataset()[0]
      num_parts = 2
  
      partitioner = Partitioner(data, num_parts, tmp_path)
      partitioner.generate_partition()
  
      meta_path = osp.join(tmp_path, "META.json")
      assert osp.exists(meta_path)
  
      for edge_type, num_edges in data.num_edges_dict.items():
          assert len(edge_type) == 3
          edge_name = EdgeTypeStr(edge_type)
          edge_map_path = osp.join(
              tmp_path, "edge_map", f"{edge_name}.pt"
          )
          assert osp.exists(edge_map_path)
          edge_map = torch.load(edge_map_path)
          assert edge_map.numel() == num_edges
  
      for node_type, num_nodes in data.num_nodes_dict.items():
          node_map_path = osp.join(
              tmp_path, "node_map", f"{node_type}.pt"
          )
          assert osp.exists(node_map_path)
          node_map = torch.load(node_map_path)
          assert node_map.numel() == num_nodes
  
      for pid in range(num_parts):
          graph_path = osp.join(tmp_path, f"part_{pid}", "graph.pt")
          assert osp.exists(graph_path)
          node_feats_path = osp.join(
              tmp_path, f"part_{pid}", "node_feats.pt"
          )
          assert osp.exists(node_feats_path)
          edge_feats_path = osp.join(
              tmp_path, f"part_{pid}", "edge_feats.pt"
          )
          assert osp.exists(edge_feats_path)
  
file_path: pytorch_geometric/test/distributed/test_partition.py
imports: import os.path
  import pytest
  import torch
  from torch_geometric.datasets import FakeDataset, FakeHeteroDataset
  from torch_geometric.distributed import LocalFeatureStore, LocalGraphStore, Partitioner
  from torch_geometric.typing import EdgeTypeStr
label_desc: Use image node_id 2 for calling a global function with example usage: test_partition_hetero_data(tmp_path) without return types