code: def W_split_tests(
      N, C, H, W, half_halo, rank, world_size, halo_ex, num_steps
  ):
      Wr = 8 * world_size
      Wp = ((W + Wr - 1) // Wr) * 8
  
      for i in range(4):
          div = int(pow(2, i))
          single_test(
              rank,
              world_size,
              halo_ex,
              C * div,
              H // div,
              Wp // div,
              half_halo,
              torch.float16,
              1,
              False,
              num_steps,
          )
          single_test(
              rank,
              world_size,
              halo_ex,
              C * div,
              H // div,
              Wp // div,
              half_halo,
              torch.float16,
              2,
              False,
              num_steps,
          )
          single_test(
              rank,
              world_size,
              halo_ex,
              C * div,
              H // div,
              Wp // div,
              half_halo,
              torch.float16,
              3,
              False,
              num_steps,
          )
  
file_path: apex/apex/contrib/test/peer_memory/test_peer_halo_exchange_module.py
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 4 for calling a global function with example usage: W_split_tests(N, C, H, W, half_halo, rank, world_size, halo_ex, num_steps) without return types