Source code for espnet.nets.pytorch_backend.transformer.mask

# Copyright 2019 Shigeki Karita
#  Apache 2.0  (

"""Mask module."""

import torch

[docs]def subsequent_mask(size, device="cpu", dtype=torch.bool): """Create mask for subsequent steps (size, size). :param int size: size of mask :param str device: "cpu" or "cuda" or torch.Tensor.device :param torch.dtype dtype: result dtype :rtype: torch.Tensor >>> subsequent_mask(3) [[1, 0, 0], [1, 1, 0], [1, 1, 1]] """ ret = torch.ones(size, size, device=device, dtype=dtype) return torch.tril(ret, out=ret)
[docs]def target_mask(ys_in_pad, ignore_id): """Create mask for decoder self-attention. :param torch.Tensor ys_pad: batch of padded target sequences (B, Lmax) :param int ignore_id: index of padding :param torch.dtype dtype: result dtype :rtype: torch.Tensor (B, Lmax, Lmax) """ ys_mask = ys_in_pad != ignore_id m = subsequent_mask(ys_mask.size(-1), device=ys_mask.device).unsqueeze(0) return ys_mask.unsqueeze(-2) & m