code: def sum(input: DTensor, *args, **kwargs):
      dim = None
      if len(args) > 0:
          dim = args[0]
      if "dim" in kwargs:
          dim = kwargs["dim"]
      if isinstance(dim, int):
          dim = (dim,)
      dtype = kwargs.get("dtype", None)
      if dim is None:
          if "keepdim" in kwargs:
              raise TypeError(
                  f"sum() received an invalid combination of arguments - got (Tensor, keepdim=bool), but expected one of\n"
                  "* (Tensor input)\n"
                  "* (Tensor input, tuple of ints dim, bool keepdim)\n"
                  "* (Tensor input, tuple of names dim, bool keepdim)"
              )
          out = input.shardings.map_reduce_shard(
              lambda x: torch.sum(x), torch.add
          )
          if dtype:
              out = out.type(dtype)
          return out
  
      keepdim = kwargs.get("keepdim", False)
      if input.shardings.shapes.axis not in dim:
          return DTensor(
              input.shardings.map_shard(
                  lambda x: torch.sum(
                      x, dim=dim, keepdim=keepdim, dtype=dtype
                  ),
                  shapes=input.shardings.shapes.squeeze(dim, keepdim),
              )
          )
  
      out = input.shardings.map_reduce_shard(
          lambda x: torch.sum(x, dim=dim, keepdim=keepdim), torch.add
      )
      if dtype:
          out = out.type(dtype)
      return out
  
file_path: FATE/python/fate/arch/tensor/distributed/_ops_agg.py
imports: from typing import Tuple
  import torch
  from ._tensor import DTensor, implements
label_desc: Use image node_id 1 for calling a global function with example usage: sum(input) and returns: out, out, DTensor