Function mean

imperative/python/megengine/xla/rules/reduction.py:115–124

(inp, axes=None, keepdims=False)

def mean(inp, axes=None, keepdims=False):
    axes = _normalize_reduce_axes(inp.shape, axes)
    inp_sum = sum(inp, axes, keepdims)
    inp_shape = inp.shape

    divider = 1.0
    for ax in axes:
        divider *= inp_shape[ax]

    return div(inp_sum, divider)


@register_lower_rule(mops.Reduce)