def test_attention_masks():
    """Correctness tests for the attention mask functions."""

    # ======================
    # test causal_attn_mask
    # ======================
    M_ref = torch.BoolTensor(
        [[1, 0, 0, 0, 0], [1, 1, 0, 0, 0], [1, 1, 1, 0, 0], [1, 1, 1, 1, 0], [1, 1, 1, 1, 1]]
    )
    M = causal_attn_mask(size=5, device="cpu")

    assert_equal(M_ref, M, msg="causal_attn_mask produced incorrect result")
Correctness tests for the attention mask functions.
neuron_explainer/tests/test_transformer.py