Feature request: pass positions of sequence items in AttentionLayers forward pass

In some applications it is desirable to have sequence items referenced by arbitrary positions instead of ascending ordinals (torch.arange(max_seq_len...)).

Would it be possible to add a new parameter pos to the forward pass of AttentionLayers?

I am not sure whether all positional encodings would be compatible but the idea would be something like:

from x_transformers import Decoder, Encoder

enc = Encoder(
    dim = 512,
    depth = 2,
    heads = 8,
    rotary_pos_emb = True,
)

#current behavior is
enc(torch.rand(2,5,512),pos=torch.arange(5).unsqueeze(0))

#arbitrary positions possible
enc(torch.rand(2,5,512),pos=torch.Float([[0,1,2,3,4.1],[-1.5,0,2,3.5,4]]))

[antorsae]

Yes, but for AttentionLayers (not just TransformerWrapper as in the PR).

Also, this would be more attractive if it worked with the other more performant (relative) positional embeddings/encodings (T5, rotary, etc.).

[lucidrains]

@antorsae or do you mean derive the relative positional encoding based on the absolute positions of the tokens passed in on forwards?