def Mean(x, axis=None, keepdims=True):
"""mean."""
outs = mean(x, axis, keepdims)
# remove useless mean_output
if isinstance(outs, tuple):
outs = outs[0]
if outs.op.name == "mean_output":
outs = outs.op.input_tensors[0]
return outs
def mean_ad(head, input_shape, axis, keepdims):
"""mean autodiff."""
tensor_a = tvm.placeholder(input_shape, head.dtype, "A")
tensor_b = mean.mean(tensor_a, axis, keepdims)
# remove useless mean_output
if isinstance(tensor_b, tuple):
tensor_b = tensor_b[0]
if tensor_b.op.name == "mean_output":
tensor_b = tensor_b.op.input_tensors[0]
jacs = list(akg.differentiate(tensor_b, [tensor_a], head))
return jacs[0]
def mean_ad(head, input_shape, axis, keepdims):
"""
Compute gradient of mean operator using automatic differentiate.
Args:
head (tvm.tensor.Tensor): Input tensor.
input_shape (Union[list, tuple]): Shape of input tensor of mean operator.
axis (Union[list, tuple, int]): Specifies which axis to reduce.
keepdims (bool): Keep the reduced axis with length 1 if keepdims is true.
Returns:
tvm.tensor.Tensor.
"""
a = akg.tvm.placeholder(input_shape, head.dtype, "A")
b, _ = mean.mean(a, axis, keepdims)
jacs = list(akg.differentiate(b, [a], head))
return jacs[0]
def mean_square(inputs, axis=None, keepdims=False):
"""Mean of square value of a tensor, alongside the specified axis.
Arguments:
input: A tensor.
axis: A list of integer. Axes to compute the mean.
keepdims: A boolean, whether to keep the dimensions or not.
If `keepdims` is `False`, the rank of the tensor is reduced
by 1 for each entry in `axis`. If `keepdims` is `True`,
the reduced dimensions are retained with length 1.
Returns:
A tensor with the mean of element-wise square value of `input`.
Notice: There is some precision problem for the operator and remain to solve
"""
inputs_square = square(inputs)
return mean(inputs_square, axis, keepdims)
def mean_mul(first_input, second_input, axis=None, keepdims=False):
temp, _ = mean.mean(first_input, axis, keepdims)
output = mul.mul(temp, second_input)
return output
def SimpleMean(x):
"""mean"""
return mean.mean(x, axis=[2, 3], keepdims=True)
def mul_mean(first_input, second_input, axis=None, keepdims=False):
temp = mul.mul(first_input, second_input)
output, _ = mean.mean(temp, axis, keepdims)
return output
def ReduceMean(x, axis, keepdims):
"""reduce_mean"""
return mean.mean(x, axis=axis, keepdims=keepdims)