Masked Softmax PyTorch

masked_softmax_pytorch.py

import torch
from torch.nn.functional import softmax

def masked_softmax(input: torch.Tensor, bool_mask: torch.Tensor, dim: int = -1, dtype: torch.dtype = None) -> torch.Tensor:
    min_type_value = torch.finfo(input.dtype).min
    masked_value = input.masked_fill(bool_mask, min_type_value)

    return softmax(masked_value, dim = dim, dtype = dtype)

## Example 
# a = torch.rand(1, 3)
# print(a) ## tensor([[0.1218, 0.7097, 0.4001]])

# b = masked_softmax(a, a < 0.5)
# print(b) ## tensor([[0., 1., 0.]])

# c = masked_softmax(a, torch.tensor([True, False, False]))
# print(c) ## tensor([[0.0000, 0.5768, 0.4232]])