code: def create_hetero_data(
      tmp_path: str,
      rank: int,
      time_attr: Optional[str] = None,
  ):
      graph_store = LocalGraphStore.from_partition(tmp_path, pid=rank)
      feature_store = LocalFeatureStore.from_partition(
          tmp_path, pid=rank
      )
      (
          meta,
          num_partitions,
          partition_idx,
          node_pb,
          edge_pb,
      ) = load_partition_info(tmp_path, rank)
      graph_store.partition_idx = (
          feature_store.partition_idx
      ) = partition_idx
      graph_store.num_partitions = (
          feature_store.num_partitions
      ) = num_partitions
      graph_store.node_pb = feature_store.node_feat_pb = node_pb
      graph_store.edge_pb = feature_store.edge_feat_pb = edge_pb
      graph_store.meta = feature_store.meta = meta
  
      if time_attr == "time":  # Create node-level time data:
          feature_store.put_tensor(
              tensor=torch.full(
                  (len(node_pb["v0"]),), 1, dtype=torch.int64
              ),
              group_name="v0",
              attr_name=time_attr,
          )
          feature_store.put_tensor(
              tensor=torch.full(
                  (len(node_pb["v1"]),), 2, dtype=torch.int64
              ),
              group_name="v1",
              attr_name=time_attr,
          )
      elif time_attr == "edge_time":  # Create edge-level time data:
          i = 0
          for attr, edge_index in graph_store._edge_index.items():
              time = torch.full(
                  (edge_index.size(1),), i, dtype=torch.int64
              )
              feature_store.put_tensor(
                  tensor=time,
                  group_name=attr[0],
                  attr_name=time_attr,
              )
              i += 1
  
      return feature_store, graph_store
  
imports: import atexit
  import socket
  from typing import Optional
  import pytest
  import torch
  from torch_geometric.data import Data
  from torch_geometric.datasets import FakeHeteroDataset
  from torch_geometric.distributed import DistNeighborSampler, LocalFeatureStore, LocalGraphStore, Partitioner
  from torch_geometric.distributed.dist_context import DistContext
  from torch_geometric.distributed.event_loop import ConcurrentEventLoop
  from torch_geometric.distributed.partition import load_partition_info
  from torch_geometric.distributed.rpc import init_rpc, shutdown_rpc
  from torch_geometric.sampler import NeighborSampler, NodeSamplerInput
  from torch_geometric.sampler.neighbor_sampler import node_sample
  from torch_geometric.testing import onlyLinux, withPackage
label_desc: Use image node_id 2 for calling a global function with example usage: create_hetero_data(tmp_path, rank, time_attr) and returns: feature_store,  graph_store