cerebras.modelzoo.common.utils.model.transformer_utils.create_2D_full_mask#

cerebras.modelzoo.common.utils.model.transformer_utils.create_2D_full_mask(src_sequence_length: int, target_sequence_length: int, dtype=None, device=None)[source]#

Create autoregressive (triangular) mask.

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].

previous

cerebras.modelzoo.common.utils.model.transformer_utils.create_2D_autoregressive_mask

next

cerebras.modelzoo.common.utils.model.transformer_utils.create_broadcasted_autoregressive_mask