code: def create_data(
      rank: int, world_size: int, time_attr: Optional[str] = None
  ):
      if rank == 0:  # Partition 0:
          node_id = torch.tensor([0, 1, 2, 3, 4, 5, 9])
          edge_index = torch.tensor(
              [  # Sorted by destination.
                  [1, 2, 3, 4, 5, 0, 0],
                  [0, 1, 2, 3, 4, 4, 9],
              ]
          )
      else:  # Partition 1:
          node_id = torch.tensor([0, 4, 5, 6, 7, 8, 9])
          edge_index = torch.tensor(
              [  # Sorted by destination.
                  [5, 6, 7, 8, 9, 5, 0],
                  [4, 5, 6, 7, 8, 9, 9],
              ]
          )
      feature_store = LocalFeatureStore.from_data(node_id)
      graph_store = LocalGraphStore.from_data(
          edge_id=None,
          edge_index=edge_index,
          num_nodes=10,
          is_sorted=True,
      )
  
      graph_store.node_pb = torch.tensor([0, 0, 0, 0, 0, 1, 1, 1, 1, 1])
      graph_store.meta.update({"num_parts": 2})
      graph_store.partition_idx = rank
      graph_store.num_partitions = world_size
  
      edge_index = torch.tensor(
          [  # Create reference data:
              [1, 2, 3, 4, 5, 0, 5, 6, 7, 8, 9, 0],
              [0, 1, 2, 3, 4, 4, 9, 5, 6, 7, 8, 9],
          ]
      )
      data = Data(x=None, y=None, edge_index=edge_index, num_nodes=10)
  
      if time_attr == "time":  # Create node-level time data:
          data.time = torch.tensor([5, 0, 1, 3, 3, 4, 4, 4, 4, 4])
          feature_store.put_tensor(
              data.time, group_name=None, attr_name=time_attr
          )
  
      elif time_attr == "edge_time":  # Create edge-level time data:
          data.edge_time = torch.tensor(
              [0, 1, 2, 3, 4, 5, 7, 7, 7, 7, 7, 11]
          )
  
          if rank == 0:
              edge_time = torch.tensor([0, 1, 2, 3, 4, 5, 11])
          if rank == 1:
              edge_time = torch.tensor([4, 7, 7, 7, 7, 7, 11])
  
          feature_store.put_tensor(
              edge_time, group_name=None, attr_name=time_attr
          )
  
      return (feature_store, graph_store), data
  
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 1 for calling a global function with example usage: create_data(rank, world_size, time_attr) and returns: feature_store,  graph_store,  data