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)