tensorplay.masked.softmax

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

Returns softmax of all the slices in the input tensor along dim while the input elements are masked out according to the boolean tensor mask.

Let x be a sequence of unmasked elements of one-dimensional slice of the input tensor. Softmax of i-th element in x is defined as exp(x[i])/sum(exp(x)).

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 softmax computation, otherwise the element is ignored.

The values of masked-out elements of the output tensor have undefined value: it may or may not be set to zero or nan; the choice may correspond to the value that leads to the most efficient storage of output tensor.

The mask of the softmax output tensor can be computed as tensorplay.broadcast_to(mask, input.shape).

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.softmax(input, 1, mask=mask)
tensor([[0.1192, 0., 0.8808],
        [-nan, -nan, -nan]])