class_docstr: AllGather op with gradient backward
class_docstr_tok: ["AllGather","op","with","gradient","backward"]
code: def backward(ctx, *grads):
      grad_list = list(grads)
      rank = distributed.get_rank()
      grad_out = grad_list[rank]
  
      dist_ops = [
          distributed.reduce(
              grad_out, rank, distributed.ReduceOp.SUM, async_op=True
          )
          if i == rank
          else distributed.reduce(
              grad_list[i], i, distributed.ReduceOp.SUM, async_op=True
          )
          for i in range(distributed.get_world_size())
      ]
      for _op in dist_ops:
          _op.wait()
  
      grad_out *= len(
          grad_list
      )  # cooperate with distributed loss function
      return (grad_out, *[None for _ in range(len(grad_list))])
  
imports: import math
  from typing import Callable
  import torch
  from torch import distributed
  from torch.nn.functional import linear, normalize
label_desc: Use image node_id 2 for calling the AllGatherFunc obj's underlying member method code with example usage: obj.backward(ctx) and returns: grad_out,  unknown