36101ab4095065a4196ff4f6437e94f0d91df4e9,official/nlp/modeling/layers/attention_test.py,CachedAttentionTest,test_masked_attention,#CachedAttentionTest#,201

Before Change

    // one element.
    mask_data = np.random.randint(
        2, size=(batch_size, from_seq_length, from_seq_length))
    masked_output_data, cache = layer([from_data, from_data], mask_data, cache)
    self.assertEqual(masked_output_data.shape, (3, 4, 8))
    self.assertEqual(cache["value"].shape, (3, 4, 2, 2))

    // Tests inputs without cache.
    masked_output_data, cache = layer([from_data, from_data, mask_data])
    self.assertEqual(masked_output_data.shape, (3, 4, 8))
    self.assertIsNone(cache)

  def test_padded_decode(self):

After Change

    layer = attention.CachedAttention(num_heads=num_heads, key_size=head_size)

    // Generate data for the input (non-mask) tensors.
    from_data = tf.zeros((batch_size, from_seq_length, 8), dtype=np.float32)
    // Invoke the data with a random set of mask data. This should mask at least
    // one element.
    mask_data = np.random.randint(
        2, size=(batch_size, from_seq_length, from_seq_length))
    masked_output_data, cache = layer(
        query=from_data, value=from_data, attention_mask=mask_data, cache=cache)
    self.assertEqual(masked_output_data.shape, (3, 4, 8))
    self.assertEqual(cache["value"].shape, (3, 4, 2, 2))

    // Tests inputs without cache.
    masked_output_data, cache = layer(
        query=from_data, value=from_data, attention_mask=mask_data)
    self.assertEqual(masked_output_data.shape, (3, 4, 8))
    self.assertIsNone(cache)

  def test_padded_decode(self):

Instances