Using `TFGPT2LMHeadModel.generate` with tf.distribute.TPUStrategy and tf.function
Author: JulesGM
Pull request: no

## High-level description
I was wondering what were the steps that one is supposed to use to be able to do call `strategy.run` on `model.generate` with tf.distribute.TPUStrategy.
Right now, I get an error message telling me that TPUStrategy doesn't support pure eager functions, so I tried to put `model.__call__` and `model.generate` in a `tf.function`, in such a way:
```python
model.__call__ = tf.function(model.__call__,     experimental_relax_shapes=True)
model.generate = tf.function(model.generate,     experimental_relax_shapes=True)
```
but it doesn't work, when trying to jit either of just one of these or both (or neither).

## Environment info
- `transformers` version: 4.5.1
- Platform: Ubuntu 20.04 LTS
- Python version: 3.8.5
- PyTorch version (GPU?): N/A
- Tensorflow version (GPU?): Tensorflow 2.5.0rc2
- Using GPU in script?: Using TPUs
- Using distributed or parallel set-up in script?: tf.distribute.TPUStrategy with 8 TPUs (single TPU host)

### Who can help
@Rocketknight1, maybe @LysandreJik 

## To reproduce

Steps to reproduce the behavior, run this very small script:

```python
import sys
import transformers
import tensorflow as tf

TPU_NAME = "" # This is dependant on your setup. This is not the problem. Works perfectly with other things.
MODEL_KEY = "distilgpt2"
EXAMPLE_SENTENCE = "I like pizza, because it's "

cr = tf.distribute.cluster_resolver.TPUClusterResolver.connect(tpu=TPU_NAME)
strategy = tf.distribute.TPUStrategy(cr)
with strategy.scope():
    model = transformers.TFGPT2LMHeadModel.from_pretrained(MODEL_KEY)
tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_KEY)
tokens = tokenizer.encode(EXAMPLE_SENTENCE, return_tensors="tf") 

# This doesn't work (See *Error Message 1*, in the next section)
strategy.run(model.generate, args=(tokens,))

# So we try to jit the __call__ function:
model.__call__ = tf.function(
    model.__call__,
    experimental_relax_shapes=True,
)

# The following also doesn't work (See *Error Message 2*, in the next section)
strategy.run(model.generate, args=(tokens,))

# We try jitting generate
model.generate = tf.function(
    model.generate,
    experimental_relax_shapes=True,
)

# Once more, it doesn't work (See *Error Message 3*, in the next section)
# Only jitting `model.generate` gives the same error.
strategy.run(model.generate, args=(tokens,))
```

## Error messages:
### Error message 1:
No method is jitted.
```
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-9-44e0c5fc714c> in <module>
----> 1 strategy.run(model.generate, args=(tokens,))

~/.local/lib/python3.8/site-packages/tensorflow/python/distribute/tpu_strategy.py in run(self, fn, args, kwargs, options)
    392       objects, or `Tensor`s (for example, if running on a single replica).
    393     """
--> 394     validate_run_function(fn)
    395 
    396     fn, args, kwargs = _maybe_partial_apply_variables(fn, args, kwargs)

~/.local/lib/python3.8/site-packages/tensorflow/python/distribute/tpu_strategy.py in validate_run_function(fn)
     98       and not isinstance(fn, function.ConcreteFunction) \
     99       and not (callable(fn) and isinstance(fn.__call__, def_function.Function)):
--> 100     raise NotImplementedError(
    101         "TPUStrategy.run(fn, ...) does not support pure eager "
    102         "execution. please make sure the function passed into "

NotImplementedError: TPUStrategy.run(fn, ...) does not support pure eager execution. please make sure the function passed into `strategy.run` is a `tf.function` or `strategy.run` is called inside a `tf.function` if eager behavior is enabled.
```

### Error message 2:
Only `model.__call__` is jitted. Same error as #1
```
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-9-44e0c5fc714c> in <module>
----> 1 strategy.run(model.generate, args=(tokens,))

~/.local/lib/python3.8/site-packages/tensorflow/python/distribute/tpu_strategy.py in run(self, fn, args, kwargs, options)
    392       objects, or `Tensor`s (for example, if running on a single replica).
    393     """
--> 394     validate_run_function(fn)
    395 
    396     fn, args, kwargs = _maybe_partial_apply_variables(fn, args, kwargs)

~/.local/lib/python3.8/site-packages/tensorflow/python/distribute/tpu_strategy.py in validate_run_function(fn)
     98       and not isinstance(fn, function.ConcreteFunction) \
     99       and not (callable(fn) and isinstance(fn.__call__, def_function.Function)):
--> 100     raise NotImplementedError(
    101         "TPUStrategy.run(fn, ...) does not support pure eager "
    102         "execution. please make sure the function passed into "

NotImplementedError: TPUStrategy.run(fn, ...) does not support pure eager execution. please make sure the function passed into `strategy.run` is a `tf.function` or `strategy.run` is called inside a `tf.function` if eager behavior is enabled.
```

### Error message 3:
This is when either just `model.generate` is jitted, or both are jitted.
```
WARNING:tensorflow:AutoGraph could not transform <bound method Socket.send of <zmq.sugar.socket.Socket object at 0x7fa6913fd640>> and will run it as-is.
Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.
Cause: module, class, method, function, traceback, frame, or code object was expected, got cython_function_or_method
To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert
WARNING: AutoGraph could not transform <bound method Socket.send of <zmq.sugar.socket.Socket object at 0x7fa6913fd640>> and will run it as-is.
Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.
Cause: module, class, method, function, traceback, frame, or code object was expected, got cython_function_or_method
To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert
Setting `pad_token_id` to 50256 (first `eos_token_id`) to generate sequence
WARNING:tensorflow:The parameters `output_attentions`, `output_hidden_states` and `use_cache` cannot be updated when calling a model.They have to be set to True/False in the config object (i.e.: `config=XConfig.from_pretrained('name', output_attentions=True)`).
WARNING:tensorflow:The parameter `return_dict` cannot be set in graph mode and will always be set to `True`.
WARNING:tensorflow:From /home/jules/.local/lib/python3.8/site-packages/tensorflow/python/ops/array_ops.py:5043: calling gather (from tensorflow.python.ops.array_ops) with validate_indices is deprecated and will be removed in a future version.
Instructions for updating:
The `validate_indices` argument has no effect. Indices are always validated on CPU and never validated on GPU.
---------------------------------------------------------------------------
OperatorNotAllowedInGraphError            Traceback (most recent call last)
<ipython-input-9-44e0c5fc714c> in <module>
----> 1 strategy.run(model.generate, args=(tokens,))

~/.local/lib/python3.8/site-packages/tensorflow/python/distribute/tpu_strategy.py in run(self, fn, args, kwargs, options)
    400     fn = autograph.tf_convert(fn, autograph_ctx.control_status_ctx())
    401     options = options or distribute_lib.RunOptions()
--> 402     return self.extended.tpu_run(fn, args, kwargs, options)
    403 
    404   def experimental_assign_to_logical_device(self, tensor, logical_device_id):

~/.local/lib/python3.8/site-packages/tensorflow/python/distribute/tpu_strategy.py in tpu_run(self, fn, args, kwargs, options)
   1426   def tpu_run(self, fn, args, kwargs, options=None):
   1427     func = self._tpu_function_creator(fn, options)
-> 1428     return func(args, kwargs)
   1429 
   1430   def _tpu_function_creator(self, fn, options):

~/.local/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
    887 
    888       with OptionalXlaContext(self._jit_compile):
--> 889         result = self._call(*args, **kwds)
    890 
    891       new_tracing_count = self.experimental_get_tracing_count()

~/.local/lib/python3.8/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

~/.local/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in _initialize(self, args, kwds, add_initializers_to)
    761     self._graph_deleter = FunctionDeleter(self._lifted_initializer_graph)
    762     self._concrete_stateful_fn = (
--> 763         self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
    764             *args, **kwds))
    765 

~/.local/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _get_concrete_function_internal_garbage_collected(self, *args, **kwargs)
   3048       args, kwargs = None, None
   3049     with self._lock:
-> 3050       graph_function, _ = self._maybe_define_function(args, kwargs)
   3051     return graph_function
   3052 

~/.local/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
   3442 
   3443           self._function_cache.missed.add(call_context_key)
-> 3444           graph_function = self._create_graph_function(args, kwargs)
   3445           self._function_cache.primary[cache_key] = graph_function
   3446 

~/.local/lib/python3.8/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   3277     arg_names = base_arg_names + missing_arg_names
   3278     graph_function = ConcreteFunction(
-> 3279         func_graph_module.func_graph_from_py_func(
   3280             self._name,
   3281             self._python_function,

~/.local/lib/python3.8/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)
    997         _, original_func = tf_decorator.unwrap(python_func)
    998 
--> 999       func_outputs = python_func(*func_args, **func_kwargs)
   1000 
   1001       # invariant: `func_outputs` contains only Tensors, CompositeTensors,

~/.local/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
    670         # the function a weak reference to itself to avoid a reference cycle.
    671         with OptionalXlaContext(compile_with_xla):
--> 672           out = weak_wrapped_fn().__wrapped__(*args, **kwds)
    673         return out
    674 

~/.local/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py in wrapper(*args, **kwargs)
    984           except Exception as e:  # pylint:disable=broad-except
    985             if hasattr(e, "ag_error_metadata"):
--> 986               raise e.ag_error_metadata.to_exception(e)
    987             else:
    988               raise

OperatorNotAllowedInGraphError: in user code:

    /home/jules/.local/lib/python3.8/site-packages/tensorflow/python/distribute/tpu_strategy.py:1448 replicated_fn  *
        result[0] = fn(*replica_args, **replica_kwargs)
    /home/jules/.local/lib/python3.8/site-packages/transformers/generation_tf_utils.py:399 generate  *
        output = self._generate_no_beam_search(
    /home/jules/.local/lib/python3.8/site-packages/transformers/generation_tf_utils.py:455 _generate_no_beam_search  *
        while cur_len < max_length:
    /home/jules/.local/lib/python3.8/site-packages/tensorflow/python/autograph/operators/control_flow.py:858 while_stmt
        _py_while_stmt(test, body, get_state, set_state, opts)
    /home/jules/.local/lib/python3.8/site-packages/tensorflow/python/autograph/operators/control_flow.py:952 _py_while_stmt
        while test():
    /home/jules/.local/lib/python3.8/site-packages/tensorflow/python/framework/ops.py:900 __bool__
        self._disallow_bool_casting()
    /home/jules/.local/lib/python3.8/site-packages/tensorflow/python/framework/ops.py:503 _disallow_bool_casting
        self._disallow_when_autograph_enabled(
    /home/jules/.local/lib/python3.8/site-packages/tensorflow/python/framework/ops.py:489 _disallow_when_autograph_enabled
        raise errors.OperatorNotAllowedInGraphError(

    OperatorNotAllowedInGraphError: using a `tf.Tensor` as a Python `bool` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.
```

Comments:
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.