+def add_memex_v3(batches, memex_proba, marker_token):
+ for input in batches:
+ if torch.rand(1).item() < memex_proba:
+ memex_len = input.size(1) // 4
+
+ t = torch.arange(input.size(1) + memex_len, device=input.device)[
+ None, :
+ ].expand(input.size(0), -1)
+
+ # Call me the tensor-spaghetti master
+
+ trigger = torch.rand(t.size(), device=t.device)
+ trigger[:, -memex_len:] = 1.0
+ trigger = (trigger.sort(dim=1).indices == 0).long()
+ memex_mask = trigger.clone()
+ memex_mask[:, memex_len:] -= memex_mask[:, :-memex_len]
+ memex_mask = memex_mask.cumsum(dim=1)
+ u = 1 - memex_mask
+ u[:, 0] = 0
+ u = u.cumsum(dim=1)
+ # assert u.min() == 0
+ # assert u.max() == input.size(1) - 1
+ v = (
+ (trigger.cumsum(dim=1) - trigger).cumsum(dim=1)
+ + torch.randint(input.size(1), (input.size(0), 1), device=t.device)
+ ) * memex_mask
+ u = u * (1 - memex_mask) + v * memex_mask
+ n = torch.arange(input.size(0), device=input.device)[:, None].expand(
+ -1, t.size(1)
+ )
+ new_input = input[n, u]
+ limits = trigger.clone()
+ limits[:, memex_len - 1 :] += limits[:, : -(memex_len - 1)]
+ new_input = new_input * (1 - limits) + memex_marker * limits
+
+ yield new_input, memex_mask
+
+ else:
+ yield input
+
+