Fix TF LED/Longformer attentions computation

[jplu]

What does this PR do?

This PR fixes the test test_saved_model_with_attentions_output for TF Longformer and LED that was failing due to an issue in computing some shapes in the attentions.

All the slow tests are now passing 🎉

[patrickvonplaten]

I don't really understand this change here. The correct shape is given by attn_probs -> I don't understand why we cannot just use shape_list(attn_probs)? IMO, something like:

attn_probs = tf.where(
     masked_index,
     tf.zeros(shape_list(masked_index), dtype=attn_probs.dtype
)

should work, no?

[patrickvonplaten]

Also, it's better to make the dtype dependent on the type of attn_probs I think

[jplu]

No need here, the default dtype of tf.zeros is always float (float16 if AMP is activated, or float32 if not).

[patrickvonplaten]

Thanks for making the test pass! Could you give some more background on why
a) tf.tile should be used instead of tf.broadcast_to &
b) why we cannot simply use the shape of attn_probs since we apply the mask on attn_probs itself? So we know that shape_list(masked_index) == shape_list(attn_probs)

[jplu]

a) tf.tile should be used instead of tf.broadcast_to &

There are two reasons for this, the first one is because broadcast_to does reshape + tile, here we don't need to reshape, just tile is enough. The second reason is that broadcast_to is not compliant with ONNXRuntime.

b) why we cannot simply use the shape of attn_probs since we apply the mask on attn_probs itself? So we know that shape_list(masked_index) == shape_list(attn_probs)

This part is a bit tricky to explain. The issue here was that attn_probs was not always the same shape, if is_global_attn is True, then the shape of attn_probs is [batch_size, seq_len, self.num_heads, self.one_sided_attn_window_size * 2 + max_num_global_attn_indices + 1], while if it equals False its shape is [batch_size, seq_len, self.num_heads, self.one_sided_attn_window_size * 2 + 1]. Now, because the shape is never potentially the same during the execution when run in graph mode, the pre-computed shape for attn_probs by the TF tracing was [batch_size, seq_len, self.num_heads, variable], where variable cannot be computed. The consequence of this was that attn_probs had never the proper shape at the end and creates a conflict in the tf.where. To solve this we had to also create a mask of a fixed shape that depends on is_global_attn.

I don't know if it is clear enough or not. Don't hesitate to tell me if there is something you don't get.

[patrickvonplaten]

a) tf.tile should be used instead of tf.broadcast_to &

There are two reasons for this, the first one is because broadcast_to does reshape + tile, here we don't need to reshape, just tile is enough. The second reason is that broadcast_to is not compliant with ONNXRuntime.

b) why we cannot simply use the shape of attn_probs since we apply the mask on attn_probs itself? So we know that shape_list(masked_index) == shape_list(attn_probs)

This part is a bit tricky to explain. The issue here was that attn_probs was not always the same shape, if is_global_attn is True, then the shape of attn_probs is [batch_size, seq_len, self.num_heads, self.one_sided_attn_window_size * 2 + max_num_global_attn_indices + 1], while if it equals False its shape is [batch_size, seq_len, self.num_heads, self.one_sided_attn_window_size * 2 + 1]. Now, because the shape is never potentially the same during the execution when run in graph mode, the pre-computed shape for attn_probs by the TF tracing was [batch_size, seq_len, self.num_heads, variable], where variable cannot be computed. The consequence of this was that attn_probs had never the proper shape at the end and creates a conflict in the tf.where. To solve this we had to also create a mask of a fixed shape that depends on is_global_attn.

I don't know if it is clear enough or not. Don't hesitate to tell me if there is something you don't get.

Thanks for the explanation - just tried it out and cool to see that your change fixes the test!

[LysandreJik]

If the slow tests pass, LGTM

[jplu]

The entire list of slow tests are ok!

[jplu]

@sgugger Feel free to merge if it looks ok for you!

[sgugger]

Thanks for fixing!