Recent Budget Simulation is incorrect

[gopikrishnajha]

The simulation code for recent budget is incorrect here. Consider a recent budget 0. Applying the mask with recent budget leaves the main diagonal as it is which should not be the case.

Here is minimal code for replication.

import torch
recent_budget = 0
attn_weights =  torch.rand(4, 4)
ones = torch.ones_like(attn_weights, dtype=torch.bool)
ones = torch.tril(ones, diagonal=recent_budget)
print(ones)
ones = torch.triu(ones, diagonal=recent_budget)
print(ones)
attn_weights[~ones] = 0
print(attn_weights)

The output is this code is this:

tensor(
[[ True, False, False, False],
[ True, True, False, False],
[ True, True, True, False],
[ True, True, True, True]])
tensor(
[[ True, False, False, False],
[False, True, False, False],
[False, False, True, False],
[False, False, False, True]])
tensor(
[[0.4726, 0.0000, 0.0000, 0.0000],
[0.0000, 0.6668, 0.0000, 0.0000],
[0.0000, 0.0000, 0.4941, 0.0000],
[0.0000, 0.0000, 0.0000, 0.5652]])

With recent_budget=0, we should expect all zeros here.