2a11d3f0bad2430b81628fa6cb0a6301d099d77d,performer/fast_attention/tensorflow/fast_attention_test.py,TransformerLayersTest,test_softmax_noncausal_attention_block_output,#TransformerLayersTest#,53

Before Change

        num_random_features, dim)
    attention_block_output = fast_attention.favor_attention(
        query, key, value, kernel_transformation, False, projection_matrix)
    self.assertListEqual(attention_block_output.get_shape().as_list(),
                         [batch_size, length, num_heads, dim])

  def test_fast_attention(self):
    hidden_size = 64
    num_heads = 4

After Change

    attention_block_output = fast_attention.favor_attention(
        query, key, value, kernel_transformation, False, projection_matrix)

    query = tf.multiply(query, 1.0 / math.sqrt(float(dim)))
    attention_scores = tf.einsum("BXHD,BYHD->BXYH", query, key)
    attention_scores = tf.nn.softmax(attention_scores, axis=2)
    exact_attention_block_output = tf.einsum("BXYH,BYHD->BXHD",
                                             attention_scores, value)
    max_error = 2.0
    error = tf.math.abs(

Instances