code: def main():
      torch.use_deterministic_algorithms(True)
  
      torch.distributed.init_process_group("nccl")
      rank = torch.distributed.get_rank()
      world_size = torch.distributed.get_world_size()
      torch.cuda.set_device(rank)
  
      explicit_nhwc = True
  
      dtype = torch.float16
      N, C, H, W = 1, 64, 200, 336
      Hs = ((H + 8 * world_size - 1) // (8 * world_size)) * 8
      H = Hs * world_size
      gt_bottleneck = ground_truth_bottleneck(C, dtype, explicit_nhwc)
      gt = ground_truth(N, C, H, W, dtype, 1, gt_bottleneck)
  
      # verify that spatial bottleneck with group_size 1 produces same results as ground truth bottleneck
      spatial_bottleneck = spatial_parallel_bottleneck(
          C, dtype, explicit_nhwc, gt_bottleneck, None
      )
      bt = apply_to_different_bottleneck(gt, spatial_bottleneck)
      compare(gt, bt)
      # print_bottleneck_p_and_b(gt_bottleneck)
      # print_bottleneck_p_and_b(spatial_bottleneck)
  
      group_size = world_size
      group = rank // group_size
      ranks = [group * group_size + i for i in range(group_size)]
      rank_in_group = rank % group_size
  
      spatial_group_size = world_size
      spatial_communicator = None
  
      peer_pool = PeerMemoryPool(0, 64 * 1024 * 1024, ranks)
  
      # class HaloExchangerNoComm(HaloExchanger):
      #    def __init__(self, ranks, rank_in_group):
      # class HaloExchangerAllGather(HaloExchanger):
      #    def __init__(self, ranks, rank_in_group, comm):
      # class HaloExchangerSendRecv(HaloExchanger):
      #    def __init__(self, ranks, rank_in_group):
      # class HaloExchangerPeer(HaloExchanger):
      #    def __init__(self, ranks, rank_in_group, peer_pool, explicit_nhwc, numSM=1):
  
      # halex = HaloExchangerAllGather(ranks, rank_in_group)
      # halex = HaloExchangerSendRecv(ranks, rank_in_group)
  
      halex = HaloExchangerPeer(
          ranks, rank_in_group, peer_pool, explicit_nhwc, numSM=0
      )
      # print("halex.signals = %s" % (str(halex.signals)))
      # Make sure peer memory halo exchanger has finished initializing flags on all ranks before proceeding
      # torch.cuda.synchronize()
      # torch.distributed.barrier()
  
      bt2 = n_way_spatial(
          halex,
          gt_bottleneck,
          gt,
          explicit_nhwc,
          world_size,
          rank,
          fp32_reduce=True,
      )
      compare(gt, bt2)
  
imports: import unittest
  import torch
  from torch.testing._internal import common_utils
  from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase
label_desc: Use image node_id 16 for calling a global function with example usage: main() without return types