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

Before Change

    num_heads = 1
    dim = 4
    num_random_features = 350
    query = tf.ones([batch_size, length, num_heads, dim])
    key = tf.ones([batch_size, length, num_heads, dim])
    value = tf.ones([batch_size, length, num_heads, dim])
    kernel_transformation = fast_attention.softmax_kernel_transformation
    projection_matrix = fast_attention.create_projection_matrix(

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(
        (exact_attention_block_output - attention_block_output) /
        exact_attention_block_output)

Instances