def METHOD_NAME(self):
    attention = [
        [
            [0.9, 0.1, 0.0, 0.0, 0.0],
            [0.2, 0.1, 0.7, 0.0, 0.0],
            [0.0, 0.1, 0.1, 0.8, 0.0],
            [0.0, 0.0, 0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0, 0.0, 0.0],
        ],
        [
            [0.8, 0.1, 0.1, 0.0, 0.0],
            [0.1, 0.9, 0.0, 0.0, 0.0],
            [0.0, 0.1, 0.8, 0.1, 0.0],
            [0.1, 0.1, 0.2, 0.6, 0.0],
            [0.0, 0.1, 0.1, 0.3, 0.5],
            [0.0, 0.0, 0.0, 0.0, 0.0],
        ],
    ]
    self.assertAllClose(
        sequence_to_sequence.mask_attention(attention, [2, 3], True, True),
        [
            [
                [0.1, 0.0, 0.0, 0.0, 0.0],
                [0.1, 0.7, 0.0, 0.0, 0.0],
                [0.1, 0.1, 0.0, 0.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0],
            ],
            [
                [0.1, 0.1, 0.0, 0.0, 0.0],
                [0.9, 0.0, 0.0, 0.0, 0.0],
                [0.1, 0.8, 0.1, 0.0, 0.0],
                [0.1, 0.2, 0.6, 0.0, 0.0],
                [0.1, 0.1, 0.3, 0.0, 0.0],
                [0.0, 0.0, 0.0, 0.0, 0.0],
            ],
        ],
    )