def test_three_dim(self):
actual = dense_attention._lower_triangular_mask([1, 3, 3])
expected = np.array(
[[[True, False, False], [True, True, False], [True, True, True]]],
dtype=np.bool_)
self.assertAllEqual(expected, actual)
def test_orthogonal_shape(self):
actual = dense_attention._lower_triangular_mask([3, 2])
expected = np.array([[True, False], [True, True], [True, True]],
dtype=np.bool_)
self.assertAllEqual(expected, actual)
