code: def slicewise(self, fn, *inputs):
      
      if fn == tf.add:
        assert len(inputs) == 2
        if isinstance(inputs[0], mtf.LazyAllreduceSum):
          # sum of LazyAllreduceSum (keep delaying the allreduce)
          return inputs[0] + inputs[1]
      # convert all inputs to LaidOutTensor where possible
      inputs = mtf.convert_args_to_laid_out_tensors(inputs)
      ret = fn(*[
          x.one_slice if isinstance(x, self.LaidOutTensor) else x
          for x in inputs])
      if isinstance(ret, tuple):
        return tuple([self.LaidOutTensor([t]) for t in ret])
      else:
        return self.LaidOutTensor([ret])
docs: Execute a function in parallel on all slices.
  
      Args:
        fn: a function from tf.Tensors to tf.Tensor or a tuple of tf.Tensors.
        *inputs: a list of inputs.  Each input is either a LaidOutTensor or
          is convertible to a tf.Tensor.
      Returns:
        a LaidOutTensor, or a tuple of LaidOutTensors if fn returns a tuple.