avoiding conditional indexing in positionalencoding to avoid possibil…

[ppadjinTT]

What does this PR do?

This PR avoids the use of conditional indexing in the SpeechT5 RelativePositionalEncoding forward pass. As explained in
#42087 this type of advanced indexing can result in zero-shaped tensors which can be problematic for custom AI accelerators / torch backends.

This PR provides a simple and equivalent alternative. Instead of using:

        pos_seq[pos_seq < -self.max_length] = -self.max_length
        pos_seq[pos_seq >= self.max_length] = self.max_length - 1

These torch operations can be used:

        pos_seq = torch.where(pos_seq < -self.max_length, -self.max_length, pos_seq)
        pos_seq = torch.where(pos_seq >= self.max_length, self.max_length - 1, pos_seq)

Fixes #42087

Quick test of equivalence

This simple script can be used to verify that indeed these two expressions are same:

import torch

def method1(pos_seq, max_length):
    pos_seq = torch.where(pos_seq < -max_length, -max_length, pos_seq)
    pos_seq = torch.where(pos_seq >= max_length, max_length - 1, pos_seq)
    return pos_seq

def method2(pos_seq, max_length):
    pos_seq[pos_seq < -max_length] = -max_length
    pos_seq[pos_seq >= max_length] = max_length - 1
    return pos_seq

seq_len = 512
max_length = 160
device = 'cpu'
pos_seq = torch.arange(0, seq_len).to(device=device, dtype=torch.long)
pos_seq = pos_seq[:, None] - pos_seq[None, :]

print(torch.allclose(method1(pos_seq, max_length), method2(pos_seq, max_length)))
print(torch.max(torch.abs(method1(pos_seq, max_length) - method2(pos_seq, max_length))))

Output of the script above:

True
tensor(0)

Who can review?

@eustlb @ebezzam

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: speecht5

[eustlb]

LGTM! thanks for the PR 🤗