tuple index out of range for FlaxMBartForConditionalGeneration
Author: bhavitvyamalik
Pull request: no

## 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
```

Comments:
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)**