Making past and mems variables have batch size as their first output dimension.
# 🚀 Feature request

And the moment, the variables **past** / **mems** have the shape: 
`(2, batch_size, num_heads, sequence_length, embed_size_per_head)` and `(mem_len, batch_size, embed_size)` , respectively.

meaning that the variable `batch_size` dim is not on first position. Change the variables structure to have `batch_size` on first position.

## Motivation

This might be confusing as all other variables have the `batch_size` dim no first position. 
Being certain that the first dimension is always the `batch_size` would be very helpful for the user. 
Normally `mems` and `past` variables are just used to speed up decoding, and not too much changed by the user (or even looked at) I think, but consistency should be good anyways.

## Your contribution

Changing this for GPT2/CTRL is very straightforward (changing three lines of code), but for xlnet and transfo-xl would probably take a slightly bigger change.


Up for discussion @LysandreJik @thomwolf @julien-c <|||||>This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
