Customizable Embedder and Logit Mapper

[pradeep-pyro]

Hi @lucidrains, I use x-transformers extensively and often run into the need to use custom embedder and logit mapper classes together with the AutoregressiveWrapper class, for e.g., to implement pointer networks (https://arxiv.org/abs/1506.03134). Currently this is not possible, so I work with the lower level Decoder class and implement the wrapper myself.

In this PR, I tried to enable support for this by slightly reorganizing the TokenEmbedding class, and adding a class called LogitMapper to be used as the TransformerWrapper.to_logits layer. The idea is that we can implement custom classes conforming to these interfaces and everything should work seamlessly. These classes accept kwargs in their forward method that allows us to pass additional tensors to them.

The unit tests all pass and I added some additional ones. I hope I didn't break anything! Let me know if you're ok to merge this change, or if you have better ways of achieving this behavior. Thank you.

[lucidrains]

@pradeep-pyro this is cool, thank you!

do you think it is possible to just keep the default as LinearNoBias, without the LogitMapper indirection?

also, perhaps we can just have token_emb_kwargs: dict and to_logits_kwargs: dict in the forward method instead

[pradeep-pyro]

Thanks for the feedback!

also, perhaps we can just have token_emb_kwargs: dict and to_logits_kwargs: dict in the forward method instead

This solution is perfect, let me refactor. With this change, I think I can keep using LinearNoBias.

[pradeep-pyro]

@lucidrains I modified the code like you suggested. It is certainly more cleaner this way. Please take another look when you get the chance. Thanks!

[lucidrains]

@pradeep-pyro looks great Pradeep! graciously accept! 🙏