def benchmark_layers_core_masking_overhead(self):
layer = core.Masking()
x = array_ops.ones((1, 1))
def fn():
layer(x)
self._run(fn, 10000)
def test_implicit_mask(self):
attention_layer = dense_attention.Attention()
q = core.Masking(1.1)(np.array([[[1.1], [1]]], dtype=np.float32))
v = core.Masking(1.2)(np.array([[[1.2], [1]]], dtype=np.float32))
actual = attention_layer([q, v])
self.assertAllClose([[[0], [1]]], actual)
def test_override_mask(self):
attention_layer = dense_attention.Attention()
q = core.Masking()(np.array([[[1.1]]], dtype=np.float32))
mask = np.array([[False]], dtype=np.bool_)
actual = attention_layer([q, q], mask=[mask, mask])
self.assertAllClose([[[0]]], actual)