cerebras.modelzoo.common.utils.model.transformer_utils.create_2D_autoregressive_mask#
- cerebras.modelzoo.common.utils.model.transformer_utils.create_2D_autoregressive_mask(src_sequence_length, target_sequence_length, dtype=None, device=None)[source]#
- Creates a reverted autoregressive (upper triangular) mask where the 0s refers to the tokens
to attend to and 1s refer to the tokens that are skipped.
- Parameters
batch_size (int) – Batch size.
src_sequence_length (int) – Sequence length of the source (num query vectors).
target_sequence_length (int) – Sequence length of the target (num key vectors).
dtype (torch.dtype) – Dtype of the resulting mask.
device – (torch.device): The device of the input to the model, used for causal mask creation.
- Returns
The causal mask of shape [src_seq_len, target_seq_len].