code: def n_way_spatial(
      halex,
      gt_bottleneck,
      gt,
      explicit_nhwc,
      world_size,
      rank,
      fp32_reduce=False,
  ):
      assert explicit_nhwc, "Only tested for explicit nhwc"
  
      x, _, dy, _, _ = gt
      N, H, W, C = list(
          x.shape
      )  # Tensor is already shaped properly for n-way parallel
      dtype = x.dtype
  
      spatial_group_size = world_size
      spatial_group_rank = rank
      spatial_communicator = None
      spatial_halo_exchanger = halex
      spatial_method = 1  # 1 -> overlap halo and main conv, 2 -> wait for halo, conv on padded x
      use_delay_kernel = False
      spatial_parallel_args = (
          spatial_group_size,
          spatial_group_rank,
          spatial_communicator,
          spatial_halo_exchanger,
          spatial_method,
          use_delay_kernel,
      )
      spatial_bottleneck = spatial_parallel_bottleneck(
          C, dtype, explicit_nhwc, gt_bottleneck, spatial_parallel_args
      )
  
      with torch.no_grad():
          Hs = H // spatial_group_size
          xs = x[
              :,
              spatial_group_rank * Hs : (spatial_group_rank + 1) * Hs,
              :,
              :,
          ].clone()
          dys = dy[
              :,
              spatial_group_rank * Hs : (spatial_group_rank + 1) * Hs,
              :,
              :,
          ].clone()
          xs.requires_grad = True
  
      spatial_bottleneck = graph_it(spatial_bottleneck, xs)
      _, y, _, dgrad, wgrad = fprop_and_bprop(
          spatial_bottleneck, xs, dys
      )
  
      # gather output pieces
      for n, p in wgrad.items():
          if fp32_reduce:
              p32 = p.float()
              torch.distributed.all_reduce(p32)
              p.copy_(p32.half())
          else:
              torch.distributed.all_reduce(p)
      ys = [torch.empty_like(y) for _ in range(spatial_group_size)]
      torch.distributed.all_gather(ys, y)
      y = torch.cat(ys, dim=1)
      dgrads = [
          torch.empty_like(dgrad) for _ in range(spatial_group_size)
      ]
      torch.distributed.all_gather(dgrads, dgrad)
      dgrad = torch.cat(dgrads, dim=1)
      return x, y, dy, dgrad, wgrad
  
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 15 for calling a global function with example usage: n_way_spatial(halex, gt_bottleneck, gt, explicit_nhwc, world_size, rank, fp32_reduce) and returns: x,  y,  dy,  dgrad,  wgrad