class_docstr: Matrix multiplication for :math:`C = A^TB`
  
  Args:
      A (:class:`torch.tensor`): matrix :math:`A`.
      B (:class:`torch.tensor`): matrix :math:`B`.
      tesseract_dim (int): dimension of TESSERACT fo 2.5D parallelism.
      out_shape (:class:`torch.size`): shape of output tensor.
      row_rank (int): the rank of row.
      col_rank (int): the rank of column.
      dep_rank (int): the rank of depth.
      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","=","A^TB","`","Args",":","A","(",":","class",":","`","torch.tensor","`",")",":","matrix",":","math",":","`","A","`",".","B","(",":","class",":","`","torch.tensor","`",")",":","matrix",":","math",":","`","B","`",".","tesseract_dim","(","int",")",":","dimension","of","TESSERACT","fo","2.5D","parallelism",".","out_shape","(",":","class",":","`","torch.size","`",")",":","shape","of","output","tensor",".","row_rank","(","int",")",":","the","rank","of","row",".","col_rank","(","int",")",":","the","rank","of","column",".","dep_rank","(","int",")",":","the","rank","of","depth",".","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 forward(
      ctx: Any,
      A: Tensor,
      B: Tensor,
      tesseract_dim: int,
      out_shape: Tuple[int, ...],
      row_rank: int,
      col_rank: int,
      dep_rank: int,
      row_parallel_mode: ParallelMode,
      col_parallel_mode: ParallelMode,
      data_parallel_rank: int,
      pipeline_parallel_rank: int,
      pipeline_parallel_size: int,
      tensor_parallel_size: int,
  ):
      assert (
          A.shape[-2] == B.shape[-2]
      ), "Invalid shapes: A={}, B={} for ATB.".format(A.shape, B.shape)
  
      if ctx:
          ctx.save_for_backward(A, B)
  
      A_shape = A.shape
      A = A.reshape((-1, A_shape[-1]))
      B_shape = B.shape
      B = B.reshape((-1, B_shape[-1]))
      C_shape = (A.shape[-1], B.shape[-1])
      C = torch.empty(
          C_shape, dtype=A.dtype, device=get_current_device()
      )
  
      # use circular buffer to store the communication tensor
      # 2 is enough for all cases
      A_list = [torch.empty_like(A) for _ in range(2)]
      C_list = [torch.empty_like(C) for _ in range(2)]
  
      row_group = gpc.get_group(row_parallel_mode)
      col_group = gpc.get_group(col_parallel_mode)
  
      src_a = (
          tesseract_dim * row_rank
          + tesseract_dim**2 * dep_rank
          + data_parallel_rank
          * pipeline_parallel_size
          * tensor_parallel_size
          + pipeline_parallel_rank * tensor_parallel_size
      )
      src_c = (
          col_rank
          + tesseract_dim**2 * dep_rank
          + data_parallel_rank
          * pipeline_parallel_size
          * tensor_parallel_size
          + pipeline_parallel_rank * tensor_parallel_size
      )
  
      opa = [None] * 2
      opr = [None] * 2
  
      A_list[0].copy_(A)
      opa[0] = dist.broadcast(
          A_list[0], src=src_a, group=row_group, async_op=True
      )
      cur = 0
  
      for i in range(tesseract_dim):
          if i != tesseract_dim - 1:
              A_list[1 - cur].copy_(A)
              opa[1 - cur] = dist.broadcast(
                  A_list[1 - cur],
                  src=src_a + 1,
                  group=row_group,
                  async_op=True,
              )
  
          if opr[cur] is not None:
              opr[cur].wait()
              if i - 2 == row_rank:
                  C.copy_(C_list[cur])
  
          if opa[cur] is not None:
              opa[cur].wait()
  
          torch.matmul(A_list[cur].transpose(0, 1), B, out=C_list[cur])
          opr[cur] = dist.reduce(
              C_list[cur], dst=src_c, group=col_group, async_op=True
          )
          cur = 1 - cur
          src_a += 1
          src_c += tesseract_dim
  
      for op in opr:
          op.wait()
  
      if tesseract_dim - 2 == row_rank:
          C.copy_(C_list[cur])
      if tesseract_dim - 1 == row_rank:
          C.copy_(C_list[1 - cur])
      out = C.reshape(out_shape)
  
      if ctx:
          ctx.tesseract_dim = tesseract_dim
          ctx.row_rank = row_rank
          ctx.col_rank = col_rank
          ctx.dep_rank = dep_rank
          ctx.row_parallel_mode = row_parallel_mode
          ctx.col_parallel_mode = col_parallel_mode
          ctx.A_shape = A_shape
          ctx.B_shape = B_shape
          ctx.data_parallel_rank = data_parallel_rank
          ctx.pipeline_parallel_rank = pipeline_parallel_rank
          ctx.pipeline_parallel_size = pipeline_parallel_size
          ctx.tensor_parallel_size = tensor_parallel_size
  
      return out
  
imports: from typing import Any, 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 1 for calling the Matmul_ATB_2p5D obj's underlying member method code with example usage: obj.forward(ctx, A, B, tesseract_dim, out_shape, row_rank, col_rank, dep_rank, row_parallel_mode, col_parallel_mode, data_parallel_rank, pipeline_parallel_rank, pipeline_parallel_size, tensor_parallel_size) and returns: out