class_docstr: Matrix multiplication for :math:`C = AB`.
  
  Args:
      A (:class:`torch.tensor`): matrix :math:`A`.
      B (:class:`torch.tensor`): matrix :math:`B`.
      summa_dim (int): dimension of SUMMA fo 2D parallelism.
      out_shape (:class:`torch.size`): shape of output tensor.
      row_rank (int, optional): the rank of row, defaults to None.
      col_rank (int, optional): the rank of column, defaults to None.
      row_parallel_mode (:class:`colossalai.legacy.context.ParallelMode`): row parallel mode.
      col_parallel_mode (:class:`colossalai.legacy.context.ParallelMode`): column parallel mode.
      data_parallel_rank (int): data parallel rank.
      pipeline_parallel_rank (int): pipeline parallel rank
      pipeline_parallel_size (int): pipeline parallel size.
      tensor_parallel_size (int): tensor parallel size.
  
  Note:
      The parallel_mode should be concluded in ``ParallelMode``. More details about ``ParallelMode`` could be found
      in `parallel_mode <https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/context/parallel_mode.py>`_
class_docstr_tok: ["Matrix","multiplication","for",":","math",":","`","C","=","AB","`",".","Args",":","A","(",":","class",":","`","torch.tensor","`",")",":","matrix",":","math",":","`","A","`",".","B","(",":","class",":","`","torch.tensor","`",")",":","matrix",":","math",":","`","B","`",".","summa_dim","(","int",")",":","dimension","of","SUMMA","fo","2D","parallelism",".","out_shape","(",":","class",":","`","torch.size","`",")",":","shape","of","output","tensor",".","row_rank","(","int",",","optional",")",":","the","rank","of","row",",","defaults","to","None",".","col_rank","(","int",",","optional",")",":","the","rank","of","column",",","defaults","to","None",".","row_parallel_mode","(",":","class",":","`","colossalai.legacy.context.ParallelMode","`",")",":","row","parallel","mode",".","col_parallel_mode","(",":","class",":","`","colossalai.legacy.context.ParallelMode","`",")",":","column","parallel","mode",".","data_parallel_rank","(","int",")",":","data","parallel","rank",".","pipeline_parallel_rank","(","int",")",":","pipeline","parallel","rank","pipeline_parallel_size","(","int",")",":","pipeline","parallel","size",".","tensor_parallel_size","(","int",")",":","tensor","parallel","size",".","Note",":","The","parallel_mode","should","be","concluded","in","``","ParallelMode","``",".","More","details","about","``","ParallelMode","``","could","be","found","in","`","parallel_mode","<","https",":","\/\/github.com\/hpcaitech\/ColossalAI\/blob\/main\/colossalai\/context\/parallel_mode.py",">","`","_"]
code: def backward(ctx: Any, output_grad: Tensor) -> Tuple[Tensor, ...]:
      A, B = ctx.saved_tensors
      with torch.no_grad():
          A_grad = Matmul_ABT_2D.apply(
              output_grad,
              B,
              ctx.summa_dim,
              ctx.A_shape,
              ctx.row_rank,
              ctx.col_rank,
              ctx.row_parallel_mode,
              ctx.col_parallel_mode,
              ctx.data_parallel_rank,
              ctx.pipeline_parallel_rank,
              ctx.pipeline_parallel_size,
              ctx.tensor_parallel_size,
          )
          B_grad = Matmul_ATB_2D.apply(
              A,
              output_grad,
              ctx.summa_dim,
              ctx.B_shape,
              ctx.row_rank,
              ctx.col_rank,
              ctx.row_parallel_mode,
              ctx.col_parallel_mode,
              ctx.data_parallel_rank,
              ctx.pipeline_parallel_rank,
              ctx.pipeline_parallel_size,
              ctx.tensor_parallel_size,
          )
      return (
          A_grad,
          B_grad,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          None,
          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 Matmul_AB_2D obj's underlying member method code with example usage: obj.backward(ctx, output_grad) and returns: A_grad,  B_grad,  None,  None,  None,  None,  None,  None,  None,  None,  None,  None