def test_casual_attention_mask_with_no_memory(self): seq_length, memory_length = 3, 0 causal_attention_mask = xlnet_base._create_causal_attention_mask( seq_length=seq_length, memory_length=memory_length) expected_output = np.array([[1, 0, 0], [1, 1, 0], [1, 1, 1]]) self.assertAllClose(causal_attention_mask, expected_output)
def test_causal_attention_mask_with_same_length(self): seq_length, memory_length = 3, 2 causal_attention_mask = xlnet_base._create_causal_attention_mask( seq_length=seq_length, memory_length=memory_length, same_length=True) expected_output = np.array([[1, 1, 1, 0, 0], [0, 1, 1, 1, 0], [0, 0, 1, 1, 1]]) self.assertAllClose(causal_attention_mask, expected_output)