How to learning to generate a single tensor (2xNxN) from a batch of features?

What kind of neural network (NN) do you suggest for learning to generate a single tensor with shape (2, n, n) from a batch of features with shape (batch_size, max_length_size, representations_size)?

The input seems to be of NLP type. If the application is not large and the max_length_size is not very large, you can try an LSTM. Else, you can try a transformer. You can also squeeze the input and try a vanilla NN.