Function dummy_padding_mask

axlearn/common/test_utils.py:616–638

Builds a dummy attention mask where non-padding tokens are followed by padding tokens. Example: batch_size: 3 max_seq_len: 3 output: [[1, 1, 0], [1, 0, 0], [1, 1, 1]] Args: batch_size: Batch size. max_seq_len: Sequence length. Returns: A

(*, batch_size: int, max_seq_len: int)

# TODO(markblee): Update to take prng_key explicitly.
def dummy_padding_mask(*, batch_size: int, max_seq_len: int) -> Tensor:
    """Builds a dummy attention mask where non-padding tokens are followed by padding tokens.

    Example:
        batch_size: 3
        max_seq_len: 3
        output: [[1, 1, 0], [1, 0, 0], [1, 1, 1]]

    Args:
        batch_size: Batch size.
        max_seq_len: Sequence length.

    Returns:
        A bool attention mask of shape [batch, max_seq_len]. A value of 0 indicates a padding
        token, whereas 1 indicates non-padding. Each example has at least one non-padding,
        followed by padding up to seq_len.
    """
    lower_diag = jnp.arange(max_seq_len)
    lower_diag = lower_diag[None, :] <= lower_diag[:, None]
    input_len = jax.random.randint(
        jax.random.PRNGKey(123), shape=(batch_size,), minval=0, maxval=max_seq_len
    )
    return lower_diag[input_len].astype(jnp.bool)


# TODO(markblee): Update to take prng_key explicitly.