When I first read your code, everything seemed very clear, except the step that created the tensor that was later used as the input to the a matrix. I knew what it was trying to do (create all possible permutations of the embeddings in a big tensor) but I didn't understand how it was doing it. So, I decided to refactor that specific step into its own separate function, and clarify the process by breaking it down into steps and adding descriptive comments.
I think this might be especially useful to people who need to modify that exact piece of code (as was the case with me).
To make sure that I haven't introduced any functional changes to the code, I wrote a test that makes sure the output of the two versions are identical, but I'm not sure what is the best way to share it with you. (I ran it myself and the results are identical.)
When I first read your code, everything seemed very clear, except the step that created the tensor that was later used as the input to the a matrix. I knew what it was trying to do (create all possible permutations of the embeddings in a big tensor) but I didn't understand how it was doing it. So, I decided to refactor that specific step into its own separate function, and clarify the process by breaking it down into steps and adding descriptive comments. I think this might be especially useful to people who need to modify that exact piece of code (as was the case with me).
To make sure that I haven't introduced any functional changes to the code, I wrote a test that makes sure the output of the two versions are identical, but I'm not sure what is the best way to share it with you. (I ran it myself and the results are identical.)