class_docstr: An autograd function used to avoid storing activations for intermediate results.
  
  NOTE: Even though the mean and var are passed into this function, we do the entire
  backward, including mean and var, here. We have to calculate statistics outside
  this function in order to avoid multiple all_reduces when using checkpointing.
class_docstr_tok: ["An","autograd","function","used","to","avoid","storing","activations","for","intermediate","results",".","NOTE",":","Even","though","the","mean","and","var","are","passed","into","this","function",",","we","do","the","entire","backward",",","including","mean","and","var",",","here",".","We","have","to","calculate","statistics","outside","this","function","in","order","to","avoid","multiple","all_reduces","when","using","checkpointing","."]
code: def backward(ctx, grad_output):
      needs_input_grad = ctx.needs_input_grad[0]
      needs_weight_grad = ctx.needs_input_grad[1]
  
      grad_input = None
      grad_weight = None
      grad_bias = None
  
      input, weight, bias, mean, invstd, total_count = ctx.saved_tensors
      process_group = ctx.process_group
  
      dim = [d for d in range(input.ndim) if d != 1]
      if needs_input_grad or needs_weight_grad:
          grad_common = torch.sum(
              (input - mean) * grad_output, dim=dim, keepdim=True
          )  # common to grad_weight and grad_invstd
  
      if needs_input_grad:
          if weight is None:  # i.e. affine is False
              grad_input = invstd * grad_output
              grad_mean = -torch.sum(grad_input, dim=dim, keepdim=True)
              grad_invstd = grad_common
          else:
              grad_input = (
                  invstd * weight.reshape_as(mean)
              ) * grad_output
              grad_mean = -torch.sum(grad_input, dim=dim, keepdim=True)
              grad_invstd = grad_common * weight.reshape_as(mean)
          grad_var = -0.5 * invstd.pow(3) * grad_invstd
          grad_mean += -2 * mean * grad_var
          grad_meansqr = grad_var
          vec = torch.cat([grad_mean, grad_meansqr])
          all_reduce_handle = dist.all_reduce(
              vec, group=process_group, async_op=True
          )
  
      if needs_weight_grad:
          grad_weight = (grad_common * invstd).resize_as(weight)
          grad_bias = torch.sum(grad_output, dim=dim)
  
      if needs_input_grad:
          all_reduce_handle.wait()
          vec = (
              vec / total_count
          )  # NOTE(msb) removed '* count' here to avoid '/  count' below
          grad_mean, grad_meansqr = vec.chunk(2)
          grad_input += grad_mean  # removed '/ count'
          grad_input += input * (2 * grad_meansqr)  # removed '/ count'
  
      return (
          grad_input,
          grad_weight,
          grad_bias,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
      )
  
file_path: fairscale/fairscale/experimental/nn/sync_batchnorm.py
imports: from typing import Any, Dict, List, Optional, Tuple
  import torch
  from torch import Tensor
  import torch.distributed
  from torch.distributed import ProcessGroup
  from fairscale.internal import torch_version
  from fairscale.nn.checkpoint import is_checkpointing, is_recomputing
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_weight,  grad_bias,  None,  None,  None,  None,  None,  None,  None