code: def backward(ctx, output_grad):
      return (
          all_gather(output_grad, ctx.dim, ctx.parallel_mode),
          None,
          None,
      )
  
imports: from typing import Any, Optional, Tuple
  import torch
  import torch.distributed
  from torch import Tensor
  from torch.cuda.amp import custom_bwd, custom_fwd
  from colossalai.legacy.communication.collective import all_gather, all_reduce, reduce_scatter
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from colossalai.utils import get_current_device
label_desc: Use image node_id 2 for calling the _ReduceScatterTensor2D obj's underlying member method code with example usage: obj.backward(ctx, output_grad) and returns: all_gather,  None,  None