tuple index out of range for FlaxMBartForConditionalGeneration
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.9.0.dev0 (installed from source)
- Platform: Google colab
- Python version: 3.7.10
- Using TPU in script?: Yes
- Dependecies were installed following this colab notebook: https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/causal_language_modeling_flax.ipynb#scrollTo=Sj1mJNJa6PPS

### Who can help
@patil-suraj @patrickvonplaten 

## Information

Model I am using: FlaxMBartForConditionalGeneration

The problem arises when loading the model itself

## To reproduce

Steps to reproduce the behavior:
```
from transformers import FlaxMBartForConditionalGeneration
model = FlaxMBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-one-to-many-mmt")
```
```
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-5-f8556949d896> in <module>()
      1 from transformers import FlaxMBartForConditionalGeneration, MBart50TokenizerFast
      2 
----> 3 model = FlaxMBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-one-to-many-mmt", from_pt=True)
      4 tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-one-to-many-mmt", src_lang="en_XX")

15 frames
/usr/local/lib/python3.7/dist-packages/transformers/modeling_flax_utils.py in from_pretrained(cls, pretrained_model_name_or_path, dtype, *model_args, **kwargs)
    336 
    337         # init random models
--> 338         model = cls(config, *model_args, **model_kwargs)
    339 
    340         if from_pt:

/usr/local/lib/python3.7/dist-packages/transformers/models/mbart/modeling_flax_mbart.py in __init__(self, config, input_shape, seed, dtype, **kwargs)
    948     ):
    949         module = self.module_class(config=config, dtype=dtype, **kwargs)
--> 950         super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)
    951 
    952     def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple) -> FrozenDict:

/usr/local/lib/python3.7/dist-packages/transformers/modeling_flax_utils.py in __init__(self, config, module, input_shape, seed, dtype)
    103 
    104         # randomly initialized parameters
--> 105         random_params = self.init_weights(self.key, input_shape)
    106 
    107         # save required_params as set

/usr/local/lib/python3.7/dist-packages/transformers/models/mbart/modeling_flax_mbart.py in init_weights(self, rng, input_shape)
    973             decoder_attention_mask,
    974             position_ids,
--> 975             decoder_position_ids,
    976         )["params"]
    977 

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in init(self, rngs, method, mutable, *args, **kwargs)
    998     _, v_out = self.init_with_output(
    999         rngs, *args,
-> 1000         method=method, mutable=mutable, **kwargs)
   1001     return v_out
   1002 

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in init_with_output(self, rngs, method, mutable, *args, **kwargs)
    967       rngs = {'params': rngs}
    968     return self.apply(
--> 969         {}, *args, rngs=rngs, method=method, mutable=mutable, **kwargs)
    970 
    971   def init(self,

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in apply(self, variables, rngs, method, mutable, capture_intermediates, *args, **kwargs)
    937         method, self,
    938         mutable=mutable, capture_intermediates=capture_intermediates
--> 939     )(variables, *args, **kwargs, rngs=rngs)
    940 
    941   def init_with_output(self,

/usr/local/lib/python3.7/dist-packages/flax/core/scope.py in wrapper(variables, rngs, *args, **kwargs)
    685               **kwargs) -> Union[Any, Tuple[Any, VariableDict]]:
    686     with bind(variables, rngs=rngs, mutable=mutable).temporary() as root:
--> 687       y = fn(root, *args, **kwargs)
    688     if mutable is not False:
    689       return y, root.mutable_variables()

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in scope_fn(scope, *args, **kwargs)
   1176     _context.capture_stack.append(capture_intermediates)
   1177     try:
-> 1178       return fn(module.clone(parent=scope), *args, **kwargs)
   1179     finally:
   1180       _context.capture_stack.pop()

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in wrapped_module_method(*args, **kwargs)
    273     _context.module_stack.append(self)
    274     try:
--> 275       y = fun(self, *args, **kwargs)
    276       if _context.capture_stack:
    277         filter_fn = _context.capture_stack[-1]

/usr/local/lib/python3.7/dist-packages/transformers/models/mbart/modeling_flax_mbart.py in __call__(self, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, position_ids, decoder_position_ids, output_attentions, output_hidden_states, return_dict, deterministic)
   1310             output_hidden_states=output_hidden_states,
   1311             return_dict=return_dict,
-> 1312             deterministic=deterministic,
   1313         )
   1314 

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in wrapped_module_method(*args, **kwargs)
    273     _context.module_stack.append(self)
    274     try:
--> 275       y = fun(self, *args, **kwargs)
    276       if _context.capture_stack:
    277         filter_fn = _context.capture_stack[-1]

/usr/local/lib/python3.7/dist-packages/transformers/models/mbart/modeling_flax_mbart.py in __call__(self, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, position_ids, decoder_position_ids, output_attentions, output_hidden_states, return_dict, deterministic)
    905             output_hidden_states=output_hidden_states,
    906             return_dict=return_dict,
--> 907             deterministic=deterministic,
    908         )
    909 

/usr/local/lib/python3.7/dist-packages/flax/linen/module.py in wrapped_module_method(*args, **kwargs)
    273     _context.module_stack.append(self)
    274     try:
--> 275       y = fun(self, *args, **kwargs)
    276       if _context.capture_stack:
    277         filter_fn = _context.capture_stack[-1]

/usr/local/lib/python3.7/dist-packages/transformers/models/mbart/modeling_flax_mbart.py in __call__(self, input_ids, attention_mask, position_ids, output_attentions, output_hidden_states, return_dict, deterministic)
    763         )
    764 
--> 765         last_hidden_states = outputs[0]
    766         last_hidden_states = self.layer_norm(last_hidden_states)
    767 

/usr/local/lib/python3.7/dist-packages/transformers/file_utils.py in __getitem__(self, k)
   1810             return inner_dict[k]
   1811         else:
-> 1812             return self.to_tuple()[k]
   1813 
   1814     def __setattr__(self, name, value):

IndexError: tuple index out of range
```
Hi @bhavitvyamalik 

I tried this on TPU VM and it works. Where did you try it colab TPU or TPU VM ? <|||||>I tried this on Colab. Wanted to test the pipeline here before shifting it TPU VM. Is there any way to run it on google colab?<|||||>Not sure, I will try to see what's the issue with colab. But it should work just fine on TPU VM.<|||||>1. You were right! Works fine with TPU VM except for this (nothing to worry about I think):
`tcmalloc: large alloc 2444541952 bytes == 0x8f822000 @  0x7f5700b36680 0x7f5700b57824 0x5f7b11 0x648631 0x5c38e6 0x4f30e6 0x64ee88 0x505653 0x56acb6 0x568d9a 0x50b868 0x56fb87 0x568d9a 0x68cdc7 0x67e161 0x67e1df 0x4a447c 0x4a4619 0x67e829 0x4eee7b 0x6b71ed 0x7f570094d0b3 0x5f96de`

2. I was trying to run a forward pass for generating outputs using this model. Using only `eval_step` part should suffice here right? (I referred to [this](https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/causal_language_modeling_flax.ipynb#scrollTo=Sj1mJNJa6PPS) notebook here for steps) 
```
linear_decay_lr_schedule_fn = optax.linear_schedule(init_value=3e-4, end_value=0, transition_steps=1000)
adamw = optax.adamw(learning_rate=linear_decay_lr_schedule_fn, b1=0.9, b2=0.98, eps=1e-8, weight_decay=0.01)

state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=adamw)

def eval_step(params, batch):
    generated_tokens = model.generate(**batch, params=params, train=False, forced_bos_token_id=tokenizer.lang_code_to_id["fr_XX"])["sequences"]
    return generated_tokens

parallel_eval_step = jax.pmap(eval_step, "batch")

for model_input in model_inputs:  # model_inputs contain tokenized values of input sentences
      output_logits = parallel_eval_step(state.params, model_input)   # Model forward
```<|||||>This [colab](https://colab.research.google.com/drive/1qn7d9FkEOEIQcaLr2WFhopH64JGKyXe6?usp=sharing) should help with how to use generate on TPU<|||||>@patil-suraj -> let's check if we can solve those issue when changing to `jnp.ndarray` type-check<|||||>fixed in #12638<|||||>I also face this error during quantization, I am using  fastt5 library to quantize the weights of  this model **"pszemraj/grammar-synthesis-base"** , but  in transformers library (one file of this library at this path(**/usr/local/lib/python3.7/dist-packages/transformers/utils/generic.py])** show the error in colab notebook.

Error is **IndexError: tuple index out of range** 

code is here:
`!pip install fastt5

from fastT5 import (OnnxT5, get_onnx_runtime_sessions,
                    generate_onnx_representation, quantize)
from transformers import AutoTokenizer

model_or_model_path = 'pszemraj/grammar-synthesis-base'
# Step 1. convert huggingfaces t5 model to onnx
onnx_model_paths = generate_onnx_representation(model_or_model_path)


# Step 2. (recommended) quantize the converted model for fast inference and to reduce model size.
quant_model_paths = quantize(onnx_model_paths)

# step 3. setup onnx runtime
model_sessions = get_onnx_runtime_sessions(quant_model_paths)

# step 4. get the onnx model
model = OnnxT5(model_or_model_path, model_sessions)`

Error occurred in this function(generate_onnx_representation)..
So how can we debug the error (thanks)..

one more error may you face during quantize is that **encoder is not defined **  at this function **(generate_onnx_representation)**