tensorplay.masked.mean

tensorplay.masked.mean(input, dim, *, keepdim=False, dtype=None, mask=None) → Tensor[source]

Returns mean of all the elements in the input tensor along the given dimension(s) dim while the input elements are masked out according to the boolean tensor mask.

By definition, the identity value of a mean operation is the mean value of the tensor. If all elements of the input tensor along given dimension(s) dim are masked-out, the identity value of the mean is undefined. Due to this ambiguity, the elements of output tensor with strided layout, that correspond to fully masked-out elements, have nan values.

If keepdim is True, the output tensor is of the same size as input except in the dimension(s) dim where it is of size 1. Otherwise, dim is squeezed (see tensorplay.squeeze()), resulting in the output tensor having 1 (or len(dim)) fewer dimension(s).

The boolean tensor mask defines the “validity” of input tensor elements: if mask element is True then the corresponding element in input tensor will be included in mean computation, otherwise the element is ignored.

When all elements of input along the given dimension dim are ignored (fully masked-out), the corresponding element of the output tensor will have undefined value: it may or may not correspond to the identity value of mean operation; the choice may correspond to the value that leads to the most efficient storage of output tensor.

The mask of the output tensor can be computed as tensorplay.any(tensorplay.broadcast_to(mask, input.shape), dim, keepdim=keepdim, dtype=tensorplay.bool).

The shapes of the mask tensor and the input tensor don’t need to match, but they must be broadcastable under the standard broadcasting rules and the dimensionality of the mask tensor must not be greater than of the input tensor.

Example:

>>> input = tensor([[-3., -2., -1.], [0., 1., 2.]])
>>> input
tensor([[-3., -2., -1.],
        [0., 1., 2.]])
>>> mask = tensor([[True, False, True], [False, False, False]], dtype=Bool)
>>> mask
tensor([[True, False, True],
        [False, False, False]], dtype=Bool)
>>> tensorplay.masked._ops.mean(input, 1, mask=mask)
tensor([-2., -nan])