code: def backward(self, grad_output):
      grad_output = grad_output.contiguous()
      saved_input, weight, mean, invstd, count_all = self.saved_tensors
      (
          need_input_grad,
          need_weight_grad,
          need_bias_grad,
      ) = self.needs_input_grad[0:3]
  
      # calculate local stats as well as grad_weight / grad_bias
      (
          sum_dy,
          sum_dy_xmu,
          grad_weight,
          grad_bias,
      ) = torch.batch_norm_backward_reduce(
          grad_output,
          saved_input,
          mean,
          invstd,
          weight,
          need_input_grad,
          need_weight_grad,
          need_bias_grad,
      )
  
      if need_input_grad:
          # synchronizing stats used to calculate input gradient.
          sum_dy_handle = allreduce_async(
              sum_dy, op=Sum, name="sync_batch_norm.sum_dy"
          )
          sum_dy_xmu_handle = allreduce_async(
              sum_dy_xmu, op=Sum, name="sync_batch_norm.sum_dy_xmu"
          )
  
          # wait on the async communication to finish
          sum_dy = synchronize(sum_dy_handle)
          sum_dy_xmu = synchronize(sum_dy_xmu_handle)
  
          if _SYNC_BN_V4:
              # from 1.9.0 on we need a count tensor on all devices
              # count_all is calculated as total count across all ranks in forward function
              count_all = count_all.to(
                  dtype=torch.int, device=grad_output.device
              )
          elif _SYNC_BN_V2 or _SYNC_BN_V3:
              # before 1.9.0 we need the count as an integer to compute means values
              count = count_all.sum()
          else:
              # before 1.5.0, sum_dy was sum of means from every worker, so we just
              # need to divide it by number of workers
              count = size()
  
          # backward pass for gradient calculation
          # we are calling into a non-public undocumented function which broke moving to 1.9.0
          # https://github.com/pytorch/pytorch/issues/57900
          if _SYNC_BN_V4:
              # from 1.9.0 on, sums and count parameters expected
              grad_input = torch.batch_norm_backward_elemt(
                  grad_output,
                  saved_input,
                  mean,
                  invstd,
                  weight,
                  sum_dy,
                  sum_dy_xmu,
                  count_all,
              )
          else:
              # before 1.9.0, mean parameters expected, not sums and count
              grad_input = torch.batch_norm_backward_elemt(
                  grad_output,
                  saved_input,
                  mean,
                  invstd,
                  weight,
                  sum_dy / count,
                  sum_dy_xmu / count,
              )
      else:
          grad_input = None
  
      # synchronizing of grad_weight / grad_bias is not needed as distributed
      # training would handle all reduce.
      if weight is None or not need_weight_grad:
          grad_weight = None
  
      if weight is None or not need_bias_grad:
          grad_bias = None
  
      return (
          grad_input,
          grad_weight,
          grad_bias,
          None,
          None,
          None,
          None,
          None,
          None,
      )
  
file_path: horovod/horovod/torch/sync_batch_norm.py
imports: from horovod.torch.mpi_ops import allgather_async, allreduce_async, Sum, size, synchronize
  from packaging import version
  import torch
  from torch.autograd.function import Function
  import torch.nn.functional
  from torch.nn.modules.batchnorm import _BatchNorm
label_desc: Use image node_id 2 for calling the _SyncBatchNorm obj's underlying member method code with example usage: obj.backward(grad_output) and returns: grad_input,  grad_weight,  grad_bias,  None,  None,  None,  None,  None,  None