code: def reduce_mean(x,
                  disable_positional_args=None,
                  output_shape=None,
                  reduced_dim=None,
                  name=None):
    
    output_shape = convert_to_shape(output_shape)
    reduced_dim = convert_to_dimension(reduced_dim)
    assert disable_positional_args is None
    output_shape = _reduction_output_shape(x, output_shape, reduced_dim)
    with tf.variable_scope(name, default_name="reduce_mean"):
      if output_shape == x.shape:
        return x
      return reduce_sum(
          x, output_shape=output_shape) * (output_shape.size / x.shape.size)
docs: Reduction on 1 or more axes.
  
    If reduced_dim is present, then only that dimension is reduced out.
    Alternatively, specify output_shape.
    Do not specify both reduced_dim and output_shape.
    If neither is specified, then all dimensions are reduced out.
  
    Args:
      x: a Tensor
      disable_positional_args: None
      output_shape: an optional Shape. Must be a subsequence of x.shape.
      reduced_dim: a mtf.Dimension
      name: an optional string
  
    Returns:
      a Tensor