Using TensorFlow XLA with MBart50 will result in a `OperatorNotAllowedInGraphError` error
Author: xhluca
Pull request: no

### System Info

I used `pip install transformers>=4.21.0` to upgrade to the latest version.

### Who can help?

@gante 

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```python
from transformers import MBart50Tokenizer, TFMBartForConditionalGeneration

model_name = "facebook/mbart-large-50-many-to-many-mmt"
model = TFMBartForConditionalGeneration.from_pretrained(model_name, from_pt=True)
tokenizer = MBart50Tokenizer.from_pretrained(model_name)

# XLA
import tensorflow as tf
xla_generate = tf.function(model.generate, jit_compile=True)

# Translation
hi_text = "जीवन एक चॉकलेट बॉक्स की तरह है।जीवन एक चॉकलेट बॉक्स की तरह है।जीवन एक चॉकलेट बॉक्स की तरह है।जीवन एक चॉकलेट बॉक्स की तरह है।जीवन एक चॉकलेट बॉक्स की तरह है।"
chinese_text = "生活就像一盒巧克力。生活就像一盒巧克力。生活就像一盒巧克力。生活就像一盒巧克力。生活就像一盒巧克力。生活就像一盒巧克力。生活就像一盒巧克力。"

tokenizer.src_lang = "hi_IN"
encoded_hi = tokenizer([hi_text]*32, padding=True, return_tensors="tf")

tokenizer.src_lang = "zh_CN"
encoded_zh = tokenizer([chinese_text]*32, padding=True, return_tensors="tf")

# translate Hindi to French
generated_tokens = xla_generate(**encoded_hi, forced_bos_token_id=tokenizer.lang_code_to_id["fr_XX"])
x = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)

# translate Chinese to English
generated_tokens = xla_generate(**encoded_zh, forced_bos_token_id=tokenizer.lang_code_to_id["en_XX"])
y = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
```

It will result in the following error message:

```
---------------------------------------------------------------------------
OperatorNotAllowedInGraphError            Traceback (most recent call last)
<timed exec> in <module>

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
    883 
    884       with OptionalXlaContext(self._jit_compile):
--> 885         result = self._call(*args, **kwds)
    886 
    887       new_tracing_count = self.experimental_get_tracing_count()

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in _call(self, *args, **kwds)
    931       # This is the first call of __call__, so we have to initialize.
    932       initializers = []
--> 933       self._initialize(args, kwds, add_initializers_to=initializers)
    934     finally:
    935       # At this point we know that the initialization is complete (or less

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in _initialize(self, args, kwds, add_initializers_to)
    758     self._concrete_stateful_fn = (
    759         self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
--> 760             *args, **kwds))
    761 
    762     def invalid_creator_scope(*unused_args, **unused_kwds):

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _get_concrete_function_internal_garbage_collected(self, *args, **kwargs)
   3064       args, kwargs = None, None
   3065     with self._lock:
-> 3066       graph_function, _ = self._maybe_define_function(args, kwargs)
   3067     return graph_function
   3068 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
   3461 
   3462           self._function_cache.missed.add(call_context_key)
-> 3463           graph_function = self._create_graph_function(args, kwargs)
   3464           self._function_cache.primary[cache_key] = graph_function
   3465 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   3306             arg_names=arg_names,
   3307             override_flat_arg_shapes=override_flat_arg_shapes,
-> 3308             capture_by_value=self._capture_by_value),
   3309         self._function_attributes,
   3310         function_spec=self.function_spec,

/opt/conda/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes, acd_record_initial_resource_uses)
   1005         _, original_func = tf_decorator.unwrap(python_func)
   1006 
-> 1007       func_outputs = python_func(*func_args, **func_kwargs)
   1008 
   1009       # invariant: `func_outputs` contains only Tensors, CompositeTensors,

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
    666         # the function a weak reference to itself to avoid a reference cycle.
    667         with OptionalXlaContext(compile_with_xla):
--> 668           out = weak_wrapped_fn().__wrapped__(*args, **kwds)
    669         return out
    670 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py in wrapper(*args, **kwargs)
    992           except Exception as e:  # pylint:disable=broad-except
    993             if hasattr(e, "ag_error_metadata"):
--> 994               raise e.ag_error_metadata.to_exception(e)
    995             else:
    996               raise

OperatorNotAllowedInGraphError: in user code:

    /opt/conda/lib/python3.7/site-packages/transformers/generation_tf_utils.py:590 generate  *
        seed=model_kwargs.pop("seed", None),
    /opt/conda/lib/python3.7/site-packages/transformers/generation_tf_utils.py:1641 _generate  *
        input_ids,
    /opt/conda/lib/python3.7/site-packages/transformers/generation_tf_utils.py:2709 beam_search_body_fn  *
        model_outputs = self(
    /opt/conda/lib/python3.7/site-packages/transformers/modeling_tf_utils.py:703 run_call_with_unpacked_inputs  *
        return func(self, **unpacked_inputs)
    /opt/conda/lib/python3.7/site-packages/transformers/models/mbart/modeling_tf_mbart.py:1328 call  *
        outputs = self.model(
    /opt/conda/lib/python3.7/site-packages/transformers/modeling_tf_utils.py:703 run_call_with_unpacked_inputs  *
        return func(self, **unpacked_inputs)
    /opt/conda/lib/python3.7/site-packages/transformers/models/mbart/modeling_tf_mbart.py:1129 call  *
        decoder_outputs = self.decoder(
    /opt/conda/lib/python3.7/site-packages/transformers/modeling_tf_utils.py:703 run_call_with_unpacked_inputs  *
        return func(self, **unpacked_inputs)
    /opt/conda/lib/python3.7/site-packages/transformers/models/mbart/modeling_tf_mbart.py:948 call  *
        positions = self.embed_positions(input_shape, past_key_values_length)
    /opt/conda/lib/python3.7/site-packages/transformers/models/mbart/modeling_tf_mbart.py:129 call  *
        bsz, seq_len = input_shape[:2]
    /opt/conda/lib/python3.7/site-packages/tensorflow/python/framework/ops.py:520 __iter__
        self._disallow_iteration()
    /opt/conda/lib/python3.7/site-packages/tensorflow/python/framework/ops.py:513 _disallow_iteration
        self._disallow_when_autograph_enabled("iterating over `tf.Tensor`")
    /opt/conda/lib/python3.7/site-packages/tensorflow/python/framework/ops.py:491 _disallow_when_autograph_enabled
        " indicate you are trying to use an unsupported feature.".format(task))

    OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.
```

### Expected behavior

This blog post shows exactly the same way to use it: https://huggingface.co/blog/tf-xla-generate

Comments:
cc @gante 

Hi @xhluca 👋 

I was able to successfully run your example on my end. Can you try to install an updated version of `transformers` to see if it solves the problem? (`pip install -U transformers`)

Thanks it works now!