code: def dist_neighbor_sampler_hetero(
      data: FakeHeteroDataset,
      tmp_path: str,
      world_size: int,
      rank: int,
      master_port: int,
      input_type: str,
      disjoint: bool = False,
  ):
      dist_data = create_hetero_data(tmp_path, rank)
  
      current_ctx = DistContext(
          rank=rank,
          global_rank=rank,
          world_size=world_size,
          global_world_size=world_size,
          group_name="dist-sampler-test",
      )
  
      num_neighbors = [-1, -1]
      dist_sampler = DistNeighborSampler(
          data=dist_data,
          current_ctx=current_ctx,
          rpc_worker_names={},
          num_neighbors=num_neighbors,
          shuffle=False,
          disjoint=disjoint,
      )
  
      # Close RPC & worker group at exit:
      atexit.register(shutdown_rpc)
  
      init_rpc(
          current_ctx=current_ctx,
          master_addr="localhost",
          master_port=master_port,
      )
      dist_sampler.init_sampler_instance()
      dist_sampler.register_sampler_rpc()
      dist_sampler.event_loop = ConcurrentEventLoop(2)
      dist_sampler.event_loop.start_loop()
  
      # Create inputs nodes such that each belongs to a different partition
      node_pb_list = dist_data[1].node_pb[input_type].tolist()
      node_0 = node_pb_list.index(0)
      node_1 = node_pb_list.index(1)
  
      input_node = torch.tensor([node_0, node_1], dtype=torch.int64)
  
      inputs = NodeSamplerInput(
          input_id=None,
          node=input_node,
          input_type=input_type,
      )
  
      # Evaluate distributed node sample function:
      out_dist = dist_sampler.event_loop.run_task(
          coro=dist_sampler.node_sample(inputs)
      )
  
      sampler = NeighborSampler(
          data=data,
          num_neighbors=num_neighbors,
          disjoint=disjoint,
      )
  
      # Evaluate node sample function:
      out = node_sample(inputs, sampler._sample)
  
      # Compare distributed output with single machine output:
      for k in data.node_types:
          assert torch.equal(
              out_dist.node[k].sort()[0], out.node[k].sort()[0]
          )
          assert (
              out_dist.num_sampled_nodes[k] == out.num_sampled_nodes[k]
          )
          if disjoint:
              assert torch.equal(
                  out_dist.batch[k].sort()[0],
                  out.batch[k].sort()[0],
              )
  
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 5 for calling a global function with example usage: dist_neighbor_sampler_hetero(data, tmp_path, world_size, rank, master_port, input_type, disjoint) without return types