Sliding window fixes

[mattdangerw]

We have an issue, particularly on the tensorflow backend, when computing the sliding window mask during generation.

• For tf, this would affect any sequence length.
• For jax and torch, this would only affect generations longer 4096.

Before: https://colab.research.google.com/gist/mattdangerw/3d7ab7fd0f2a1169e67d3f4d43d40701/keras-tf-bug.ipynb
After: https://colab.research.google.com/gist/mattdangerw/b48e47107a2513c61d8e70a4652df468/keras-tf-bug-with-fix.ipynb

[SamanehSaadat]

Thanks for the fix, Matt!

Are these the two main bugs?

• It's been assumed that key_len==query_len.
• Caching hasn't been handled.

[grasskin]

LGTM Thank you for this fix!

[mattdangerw]

@SamanehSaadat

It's been assumed that key_len==query_len.
Caching hasn't been handled.

Kind of? Tensorflow was taking the min(query_len, sliding_window_size) as the effective sliding window size, which was basically turning the model into something that could only look 1 token behind. The general shape for generation is query_len=1, key_len=max_length.

And no backend was taking the index of generation (the cache index), to make sure our sliding window was correct for our current position.