Fix dec_attn_mask in TFTransfoXLMainLayer
Author: ydshieh
Pull request: yes

# What does this PR do?

Fix `dec_attn_mask` computation in `TFTransfoXLMainLayer`. See the section `Code Snippet to show the issue and the effect of the fix` for the details.

(This difference is more difficult for me to identify the cause).

**Some remarks**:
  - If `self.same_length` is `False`, the PT & (original) TF implementations give the same results.
  - In PT code, [self.mem_len](https://github.com/huggingface/transformers/blob/b87c044c79a408c0a1e7f7b046b5b4ce999c2d0e/src/transformers/models/transfo_xl/modeling_transfo_xl.py#L936) is used (when `self.same_length` is `True`). However, in TF code, `self.mem_len` is not used at all (inside `call`). --> First sign of inconsistency.
  - For the case `self.mlen == mlen + 1` (and `self.same_length==True`), the PT & (original) TF give the same result.
  - However, in general case, [torch.tril(all_ones, -mask_shift_len)](https://github.com/huggingface/transformers/blob/b87c044c79a408c0a1e7f7b046b5b4ce999c2d0e/src/transformers/models/transfo_xl/modeling_transfo_xl.py#L941) part in the PT code is not the same as [mask_l - mask_dia](https://github.com/huggingface/transformers/blob/b87c044c79a408c0a1e7f7b046b5b4ce999c2d0e/src/transformers/models/transfo_xl/modeling_tf_transfo_xl.py#L607) in the TF code.

## results after the fix
This inconsistency between PT / TF produces large difference in the (extended) PT/TF equivalence:
  - diffs in attention outputs are as high as `0.03`. After this fix: diffs are in the range `1e-8 ~ 1e-7`
  - diffs in hidden states are as high as `0.001`. After this fix: in the range `1e-7 ~ 2e-6`

## Code Snippet to show the issue and the effect of the fix


```
import tensorflow as tf
import torch


def pt_attn(self_mlen, mlen, qlen, same_length):
    """TransfoXLModel.forward()'s `dec_attn_mask`

    self_mlen: `self.mlen` int the original code
    """

    klen = mlen + qlen
    if same_length:
        all_ones = torch.ones((qlen, klen), dtype=torch.uint8)
        mask_len = klen - self_mlen
        if mask_len > 0:
            mask_shift_len = qlen - mask_len
        else:
            mask_shift_len = qlen
        dec_attn_mask = (torch.triu(all_ones, 1 + mlen) + torch.tril(all_ones, -mask_shift_len))  ###[:, :, None]  # -1
    else:
        dec_attn_mask = torch.triu(torch.ones((qlen, klen), dtype=torch.uint8), diagonal=1 + mlen)  ###[:, :, None]

    return dec_attn_mask.numpy()


def tf_attn(self_mlen, mlen, qlen, same_length):
    """Current TFTransfoXLMainLayer.call()'s `dec_attn_mask`

    self_mlen: `self.mlen` int the original code
    """

    attn_mask = tf.ones([qlen, qlen])
    mask_u = tf.linalg.band_part(attn_mask, 0, -1)
    mask_dia = tf.linalg.band_part(attn_mask, 0, 0)
    attn_mask_pad = tf.zeros([qlen, mlen])
    dec_attn_mask = tf.concat([attn_mask_pad, mask_u - mask_dia], 1)
    if same_length:
        mask_l = tf.linalg.band_part(attn_mask, -1, 0)
        dec_attn_mask = tf.concat([dec_attn_mask[:, :qlen] + mask_l - mask_dia, dec_attn_mask[:, qlen:]], 1)

    return tf.cast(dec_attn_mask, dtype=tf.int32).numpy()


def tf_attn_fixed(self_mlen, mlen, qlen, same_length):
    """Fixed `dec_attn_mask` in TFTransfoXLMainLayer.call()

    self_mlen: `self.mlen` int the original code
    """

    klen = mlen + qlen
    dec_attn_mask = 1 - tf.linalg.band_part(
        tf.ones([qlen, klen], dtype=tf.int32), -1, mlen
    )  # (q, q): diagonal with 1's
    if same_length:
        mask_len = klen - self_mlen
        if mask_len > 0:
            mask_shift_len = qlen - mask_len
        else:
            mask_shift_len = qlen
        if mask_shift_len >= 1:
            dec_attn_mask += 1 - tf.linalg.band_part(tf.ones([qlen, klen], dtype=tf.int32), mask_shift_len - 1, -1)
        else:
            dec_attn_mask += tf.linalg.band_part(tf.ones([qlen, klen], dtype=tf.int32), -1, -mask_shift_len)

    return dec_attn_mask.numpy()


self_mlen = 4
mlen = 4
qlen = 3
same_length = True  # False


pto = pt_attn(self_mlen, mlen, qlen, same_length)
print(pto)

tfo = tf_attn(self_mlen, mlen, qlen, same_length)
print(tfo)

tfo_fixed = tf_attn_fixed(self_mlen, mlen, qlen, same_length)
print(tfo_fixed)
```

This gives

```
# PT
[[1 0 0 0 0 1 1]
 [1 1 0 0 0 0 1]
 [1 1 1 0 0 0 0]]

# TF (diff. from PT)
[[0 0 0 0 0 1 1]
 [1 0 0 0 0 0 1]
 [1 1 0 0 0 0 0]]

# Fixed (same as PT)
[[1 0 0 0 0 1 1]
 [1 1 0 0 0 0 1]
 [1 1 1 0 0 0 0]]
```

I have further randomly generated the arguments for 1000 times, and the fixed TF version gives the same results as the PT version for all of them.

Comments:
_The documentation is not available anymore as the PR was closed or merged._

Could @gante or @Rocketknight1 give this a look?

I added a few remarks in the PR description. Maybe it will be helpful to get the ideas.

> I have further randomly generated the arguments for 1000 times, and the fixed TF version gives the same results as the PT version for all of them.

Thanks for the thorough testing <3 Can I merge the PR, @ydshieh ?

> > I have further randomly generated the arguments for 1000 times, and the fixed TF version gives the same results as the PT version for all of them.
> 
> Thanks for the thorough testing <3 Can I merge the PR, @ydshieh ?

Yes, go ahead! Thanks for reviewing.