def masked_softmax(X, valid_lens):
"""Perform softmax operation by masking elements on the last axis."""
# `X`: 3D tensor, `valid_lens`: 1D or 2D tensor
if valid_lens is None:
return nn.functional.softmax(X, dim=-1)
else:
shape = X.shape
if valid_lens.dim() == 1:
valid_lens = torch.repeat_interleave(valid_lens, shape[1])
else:
valid_lens = valid_lens.reshape(-1)
# On the last axis, replace masked elements with a very large negative
# value, whose exponentiation outputs 0
X = d2l.sequence_mask(X.reshape(-1, shape[-1]), valid_lens, value=-1e6)
return nn.functional.softmax(X.reshape(shape), dim=-1)
def masked_softmax(X, valid_len):
"""Perform softmax by filtering out some elements."""
# X: 3-D tensor, valid_len: 1-D or 2-D tensor
if valid_len is None:
return nn.functional.softmax(X, dim=-1)
else:
shape = X.shape
if valid_len.dim() == 1:
valid_len = torch.repeat_interleave(valid_len,
repeats=shape[1],
dim=0)
else:
valid_len = valid_len.reshape(-1)
# Fill masked elements with a large negative, whose exp is 0
X = d2l.sequence_mask(X.reshape(-1, shape[-1]), valid_len, value=-1e6)
return nn.functional.softmax(X.reshape(shape), dim=-1)
def masked_softmax(X, valid_len):
"""Perform softmax by filtering out some elements."""
if valid_len is None:
return nn.functional.softmax(X, dim=-1)
else:
shape = X.shape
if valid_len.dim() == 1:
# repeat value [2, 2, 2, 3, 3, 3]
valid_len = torch.repeat_interleave(valid_len,
repeats=shape[1],
dim=0)
else:
valid_len = valid_len.reshape(-1)
# print(X.reshape(-1, shape[-1]))
X = d2l.sequence_mask(X.reshape(-1, shape[-1]), valid_len, value=-1e6)
return nn.functional.softmax(X.reshape(shape), dim=-1)