code: def matmul_2d(
      a,
      b,
      summa_dim,
      out_shape,
      row_rank=None,
      col_rank=None,
      row_parallel_mode=ParallelMode.PARALLEL_2D_ROW,
      col_parallel_mode=ParallelMode.PARALLEL_2D_COL,
  ):
      r"""Matrix multiplication for 2D parallelism.
  
      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`, optional):
              row parallel mode, defaults to ParallelMode.PARALLEL_2D_ROW.
          col_parallel_mode (:class:`colossalai.legacy.context.ParallelMode`, optional):
              column parallel mode, defaults to ParallelMode.PARALLEL_2D_COL.
  
      Returns:
          :class:`torch.tensor`: :math:`C = AB`.
  
      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>`_
      """
      if row_rank is None:
          row_rank = gpc.get_local_rank(col_parallel_mode)
      if col_rank is None:
          col_rank = gpc.get_local_rank(row_parallel_mode)
  
      data_parallel_rank = (
          0
          if not gpc.is_initialized(ParallelMode.DATA)
          else gpc.get_local_rank(ParallelMode.DATA)
      )
      pipeline_parallel_rank = (
          0
          if not gpc.is_initialized(ParallelMode.PIPELINE)
          else gpc.get_local_rank(ParallelMode.PIPELINE)
      )
      pipeline_parallel_size = (
          1
          if not gpc.is_initialized(ParallelMode.PIPELINE)
          else gpc.get_world_size(ParallelMode.PIPELINE)
      )
      tensor_parallel_size = summa_dim**2
      return Matmul_AB_2D(
          a,
          b,
          summa_dim,
          out_shape,
          row_rank,
          col_rank,
          row_parallel_mode,
          col_parallel_mode,
          data_parallel_rank,
          pipeline_parallel_rank,
          pipeline_parallel_size,
          tensor_parallel_size,
      )
  
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 1 for calling a global function with example usage: matmul_2d(a, b, summa_dim, out_shape, row_rank, col_rank, row_parallel_mode, col_parallel_mode) and returns: Matmul_AB_2D