code: def matmul(a: DTensor, b: DTensor):
      assert isinstance(a, DTensor) or isinstance(
          b, DTensor
      ), "atleast one dtensor"
      if not isinstance(a, DTensor):
          return rmatmul_f(b, a)
  
      if len(a.shape) == 1 and len(b.shape) == 1:
          if isinstance(b, DTensor):
              return a.shardings.join_reduce_shard(
                  b.shardings, promote_torch_matmul, torch.add
              )
          else:
              assert (
                  a.shape[0] == b.shape[0]
              ), f"shapes mismatch: {a.shape} and {b.shape}"
              logger.warning(
                  "matmul shape 1 distributed tensor with local shape 1 tensor maybe slow"
              )
              return a.shardings.map_reduce_shard_with_stride(
                  lambda stride, size, s: promote_torch_matmul(
                      s, b[stride : stride + size]
                  ),
                  torch.add,
              )
  
      elif len(a.shape) == 1 and len(b.shape) > 1:
          if isinstance(b, DTensor):
              assert (
                  b.shardings.shapes.axis == len(b.shardings.shape) - 2
              ), "distributed axis mismatch"
              return a.shardings.join_reduce_shard(
                  b.shardings, promote_torch_matmul, torch.add
              )
          else:
              assert (
                  a.shape[0] == b.shape[-2:][0]
              ), f"shapes mismatch: {a.shape} and {b.shape}"
              logger.warning(
                  "matmul shape 1 distributed tensor with local tensor maybe slow"
              )
              axis = len(b.shape) - 2
  
              def _mapper(stride, size, s):
                  slices = tuple(
                      slice(stride, stride + size, 1)
                      if i == axis
                      else slice(None, None, None)
                      for i in range(len(b.shape))
                  )
                  return promote_torch_matmul(s, b[slices])
  
              return a.shardings.map_reduce_shard_with_stride(
                  _mapper, torch.add
              )
  
      elif len(a.shape) > 1 and len(b.shape) == 1:
          if isinstance(b, DTensor):
              assert (
                  a.shardings.shapes.axis == len(a.shardings.shape) - 1
              ), "distributed axis mismatch"
              return a.shardings.join_reduce_shard(
                  b.shardings, promote_torch_matmul, torch.add
              )
          else:
              assert (
                  a.shape[-1] == b.shape[0]
              ), f"shapes mismatch: {a.shape} and {b.shape}"
              logger.warning(
                  "matmul shape 1 distributed tensor with local tensor maybe slow"
              )
  
              def _mapper(stride, size, s):
                  slices = slice(stride, stride + size, 1)
                  return promote_torch_matmul(s, b[slices])
  
              return a.shardings.map_reduce_shard_with_stride(
                  _mapper, torch.add
              )
  
      else:
          if isinstance(b, DTensor):
              na_axis = a.shardings.shapes.axis - len(a.shape)
              nb_axis = b.shardings.shapes.axis - len(b.shape)
              # distributed axis in broadcast part
              if na_axis < -2 and nb_axis < -2 and na_axis == nb_axis:
                  shapes = [
                      torch.Size(
                          [
                              *torch.broadcast_shapes(sa[:-2], sb[:-2]),
                              sa[-2],
                              sb[-1],
                          ]
                      )
                      for sa, sb in zip(
                          a.shardings.shapes.shapes,
                          b.shardings.shapes.shapes,
                      )
                  ]
                  axis = len(shapes[0]) + na_axis
                  return DTensor(
                      a.shardings.join_shard(
                          b.shardings,
                          func=promote_torch_matmul,
                          out_shapes=shapes,
                          out_axis=axis,
                      )
                  )
  
              # distributed axis in matmul part
              elif na_axis == -1 and nb_axis == -2:
                  return a.shardings.join_reduce_shard(
                      b.shardings,
                      mapper_func=promote_torch_matmul,
                      reduce_func=torch.add,
                  )
              else:
                  raise RuntimeError(
                      f"invalid shape {a.shardings.shapes} and {b.shardings.shapes}"
                  )
  
          else:
              assert (
                  a.shape[-1] == b.shape[-2]
              ), f"shapes mismatch: {a.shape} and {b.shape}"
              na_axis = a.shardings.shapes.axis - len(a.shape)
              if na_axis != -1:
                  shapes = [
                      torch.Size(
                          [
                              *torch.broadcast_shapes(
                                  sa[:-2], b.shape[:-2]
                              ),
                              sa[-2],
                              b.shape[-1],
                          ]
                      )
                      for sa in a.shardings.shapes.shapes
                  ]
                  axis = len(shapes[0]) + na_axis
                  return DTensor(
                      a.shardings.map_shard(
                          lambda x: promote_torch_matmul(x, b),
                          shapes=shapes,
                          axis=axis,
                      )
                  )
              else:
                  logger.warning(
                      "matmul shape 1 distributed tensor with local tensor maybe slow"
                  )
                  axis = len(b.shape) - 2
  
                  def _mapper(stride, size, s):
                      slices = tuple(
                          slice(stride, stride + size, 1)
                          if i == axis
                          else slice(None, None, None)
                          for i in range(len(b.shape))
                      )
                      return promote_torch_matmul(s, b[slices])
  
                  return a.shardings.map_reduce_shard_with_stride(
                      _mapper, torch.add
                  )
  
file_path: FATE/python/fate/arch/tensor/distributed/_op_matmul.py
imports: import logging
  import torch
  from fate.arch.tensor import _custom_ops
  from ._tensor import DTensor, implements
label_desc: Use image node_id 4 for calling a global function with example usage: matmul(a, b) and returns: rmatmul_f, a, a, a, a, promote_torch_matmul, a, a, promote_torch_matmul, DTensor, DTensor, a, a, promote_torch_matmul