[Flax]Not able to Run Hugging Face GPT2 model for jax on TPU's
 Hi,Trying to do FlaxGPT2ForMultipleChoice.I am trying to run GPT2 using hugging face for TPU's.It is showing tuple out of index.But when i run it in CPU there is no such error.Even when i run a simple model without any code of mine it is also behaving the same in TPU.


![image](https://user-images.githubusercontent.com/58116635/124922271-f5a3ff00-e016-11eb-8679-3258fc3727e9.png)


This below example is for a basic model 


![image](https://user-images.githubusercontent.com/58116635/124922511-313ec900-e017-11eb-9d2c-7473fe8b21af.png)
To Reproduce
This is the colab notebook:
https://colab.research.google.com/drive/1h8CeTM5NUpHeS1oGHONX1YbwtGHf3nyU?usp=sharing

@patrickvonplaten 
<|||||>Hey @vivekvkashyap,
instead of putting a screenshot here - could you maybe copy paste the link to your google colab instead so that we can reproduce the error?

Thank you!<|||||>@patrickvonplaten i have added the  colab notebook in the section To Reproduce
<|||||>@patil-suraj
<|||||>Similar to #12578 - we are working on it :-)
