code: def mean(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 dtype is None:
          if not input.dtype.is_floating_point:
              raise RuntimeError(
                  f"mean(): could not infer output dtype. Input dtype must be either a floating point or complex dtype. Got: {input.dtype}"
              )
          dtype = input.dtype
      if dim is None:
          if "keepdim" in kwargs:
              raise TypeError(
                  f"mean() 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)"
              )
          return torch.div(
              input.shardings.map_reduce_shard(
                  lambda x: torch.sum(x, dtype=torch.float64), torch.add
              ),
              input.shape.numel(),
          ).type(dtype)
  
      keepdim = kwargs.get("keepdim", False)
      count = 1
      for d in dim:
          count *= input.shape[d]
      if input.shardings.shapes.axis not in dim:
          return DTensor(
              input.shardings.map_shard(
                  lambda x: torch.div(
                      torch.sum(
                          x,
                          dim=dim,
                          keepdim=keepdim,
                          dtype=torch.float64,
                      ),
                      count,
                  ).type(dtype),
                  shapes=input.shardings.shapes.squeeze(dim, keepdim),
              )
          )
  
      return torch.div(
          input.shardings.map_reduce_shard(
              lambda x: torch.sum(
                  x, dim=dim, keepdim=keepdim, dtype=torch.float64
              ),
              torch.add,
          ),
          count,
      ).type(dtype)
  
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 2 for calling a global function with example usage: mean(input) and returns: torch, torch, DTensor