code: def backward(ctx, grad_output):
      grad_output = grad_output.contiguous()
      # mini batch mean & var are calculated by forward path.
      # mu = 1./N*np.sum(h, axis = 0)
      # var = 1./N*np.sum((h-mu)**2, axis = 0)
      (
          saved_input,
          weight,
          mean,
          inv_std,
          z,
          bias,
          count,
      ) = ctx.saved_tensors
      process_group = ctx.process_group
      channel_last = ctx.channel_last
      world_size = ctx.world_size
      fuse_relu = ctx.fuse_relu
      grad_input = grad_z = grad_weight = grad_bias = None
  
      if fuse_relu:
          grad_output = syncbn.relu_bw_c_last(
              grad_output, saved_input, z, mean, inv_std, weight, bias
          )
      if isinstance(z, torch.Tensor) and ctx.needs_input_grad[1]:
          grad_z = grad_output.clone()
  
      # TODO: update kernel to not pre_divide by item_num
      if channel_last:
          (
              sum_dy,
              sum_dy_xmu,
              grad_weight,
              grad_bias,
          ) = syncbn.reduce_bn_c_last(
              grad_output, saved_input, mean, inv_std, weight
          )
      else:
          sum_dy, sum_dy_xmu, grad_weight, grad_bias = syncbn.reduce_bn(
              grad_output, saved_input, mean, inv_std, weight
          )
  
      # calculate grad_input
      if ctx.needs_input_grad[0]:
          if torch.distributed.is_initialized():
              num_channels = sum_dy.shape[0]
              combined = torch.cat([sum_dy, sum_dy_xmu], dim=0)
              torch.distributed.all_reduce(
                  combined,
                  torch.distributed.ReduceOp.SUM,
                  process_group,
                  async_op=False,
              )
              sum_dy, sum_dy_xmu = torch.split(combined, num_channels)
  
          if channel_last:
              grad_input = syncbn.batchnorm_backward_c_last(
                  grad_output,
                  saved_input,
                  mean,
                  inv_std,
                  weight,
                  sum_dy,
                  sum_dy_xmu,
                  count,
              )
          else:
              grad_input = syncbn.batchnorm_backward(
                  grad_output,
                  saved_input,
                  mean,
                  inv_std,
                  weight,
                  sum_dy,
                  sum_dy_xmu,
                  count,
              )
  
      if weight is None or not ctx.needs_input_grad[2]:
          grad_weight = None
  
      if weight is None or not ctx.needs_input_grad[3]:
          grad_bias = None
  
      return (
          grad_input,
          grad_z,
          grad_weight,
          grad_bias,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
      )
  
file_path: apex/apex/parallel/optimized_sync_batchnorm_kernel.py
imports: import torch
  from torch.autograd.function import Function
  import syncbn
  from apex.parallel import ReduceOp
label_desc: Use image node_id 2 for calling the SyncBatchnormFunction obj's underlying member method code with example usage: obj.backward(ctx, grad_output) and returns: grad_input,  grad_z,  grad_weight,  grad_bias,  None,  None,  None,  None,  None,  None,  None,  None