Not able use TF Dataset on TPU when created via generator in Summarization example
Author: prikmm
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.1
- Platform: Kaggle/Colab
- Python version: 3.7.10
- Tensorflow version (GPU?): 2.4.1 / 2.5.1

### Who can help
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @Rocketknight1

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

Model hub:

- for issues with a model report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->
@patil-suraj, @Rocketknight1

## Information

Model I am using (Bert, XLNet ...):

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [x] an official GLUE/SQUaD task: (XSum)
* [ ] my own task or dataset: (give details below)

I am trying to replicate the summarization example present [here](https://github.com/huggingface/transformers/blob/master/examples/tensorflow/summarization/run_summarization.py) on XSum dataset using T5, but, am facing error when trying to use a TPU (it works on gpu). 

## To reproduce

[Kaggle link](https://www.kaggle.com/rehanwild/tpu-tf-huggingface-error?scriptVersionId=69298817)

Error in TF 2.4.1:
```
---------------------------------------------------------------------------
UnavailableError                          Traceback (most recent call last)
<ipython-input-11-8513f78e8e35> in <module>
     72         model.fit(tf_tokenized_train_ds,
     73                   validation_data=tf_tokenized_valid_ds,
---> 74                   epochs=1,
     75                  )
     76                   #callbacks=[WandbCallback()])

/opt/conda/lib/python3.7/site-packages/wandb/integration/keras/keras.py in new_v2(*args, **kwargs)
    122             for cbk in cbks:
    123                 set_wandb_attrs(cbk, val_data)
--> 124         return old_v2(*args, **kwargs)
    125 
    126     training_arrays.orig_fit_loop = old_arrays

/opt/conda/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py in fit(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, max_queue_size, workers, use_multiprocessing)
   1100               tmp_logs = self.train_function(iterator)
   1101               if data_handler.should_sync:
-> 1102                 context.async_wait()
   1103               logs = tmp_logs  # No error, now safe to assign to logs.
   1104               end_step = step + data_handler.step_increment

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/context.py in async_wait()
   2328   an error state.
   2329   """
-> 2330   context().sync_executors()
   2331 
   2332 

/opt/conda/lib/python3.7/site-packages/tensorflow/python/eager/context.py in sync_executors(self)
    643     """
    644     if self._context_handle:
--> 645       pywrap_tfe.TFE_ContextSyncExecutors(self._context_handle)
    646     else:
    647       raise ValueError("Context is not initialized.")

UnavailableError: 9 root error(s) found.
  (0) Unavailable: {{function_node __inference_train_function_49588}} failed to connect to all addresses
Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
:{"created":"@1627548744.739596558","description":"Failed to pick subchannel","file":"third_party/grpc/src/core/ext/filters/client_channel/client_channel.cc","file_line":4143,"referenced_errors":[{"created":"@1627548744.739593083","description":"failed to connect to all addresses","file":"third_party/grpc/src/core/ext/filters/client_channel/lb_policy/pick_first/pick_first.cc","file_line":398,"grpc_status":14}]}
	 [[{{node MultiDeviceIteratorGetNextFromShard}}]]
	 [[RemoteCall]]
	 [[IteratorGetNextAsOptional]]
	 [[cond_14/switch_pred/_200/_88]]
  (1) Unavailable: {{function_node __inference_train_function_49588}} failed to connect to all addresses
Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
:{"created":"@1627548744.739596558","description":"Failed to pick subchannel","file":"third_party/grpc/src/core/ext/filters/client_channel/client_channel.cc","file_line":4143,"referenced_errors":[{"created":"@1627548744.739593083","description":"failed to connect to all addresses","file":"third_party/grpc/src/core/ext/filters/client_channel/lb_policy/pick_first/pick_first.cc","file_line":398,"grpc_status":14}]}
	 [[{{node MultiDeviceIteratorGetNextFromShard}}]]
	 [[RemoteCall]]
	 [[IteratorGetNextAsOptional]]
	 [[strided_slice_18/_288]]
  (2) Unavailable: {{function_node __inference_train_function_49588}} failed to connect to all addresses
Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
:{"created":"@1627548744.739596558","description":"Failed to pick subchannel","file":"third_party/grpc/src/core/ext/filters/client_channel/client_channel.cc","file_line":4143,"referenced_errors":[{"created":"@1627548744.739593083","description":"failed to connect to all addresses","file":"third_party/grpc/src/core/ext/filters/client_channel/lb_policy/pick_first/pick_first.cc","file_line":398,"grpc_status":14}]}
	 [[{{node MultiDeviceIteratorGetNextFromShard}}]]
	 [[RemoteCall]]
	 [[IteratorGetNextAsOptional]]
	 [[tpu_compile_succeeded_assert/_1965840270157496994/_8/_335]]
  (3) Unavailable: {{function_node __inference_train_function_49588}} failed to connect to all addresses
Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
:{"created":"@1627548744.739596558","description":"Failed to pick subchannel","file":"third_party/grpc/src/core/ext/filters/client_channel/client_channel.cc","file_line":4143,"referenced_errors":[{"created":"@1627548744.739593083","description":"failed to connect to all addresses","file":"third_party/grpc/src/core/ext/filters/client_channel/lb_policy/pick_first/pick_first.cc","file_line":398,"grpc_status":14}]}
	 [[{{node MultiDeviceIteratorGetNextFromShard}}]]
	 [[RemoteCall]]
	 [[IteratorGetNextAsOptional]]
	 [[Pad_27/paddings/_218]]
  (4) Unavailable:  ... [truncated]
```

Error in TF 2.5.1:
```
NotFoundError: Op type not registered 'XlaSetDynamicDimensionSize' in binary running on n-f62ff7a1-w-0. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) `tf.contrib.resampler` should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed.
```
<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

No such error
<!-- A clear and concise description of what you would expect to happen. -->

EDIT:

I found tensorflow/tensorflow#48268, though it has been closed it is not yet completely solved I guess, since I found tensorflow/tensorflow#50980. I was not able to try with TF-2.6.0-rc1 as it is not yet supported by transformers. Since, this is an upstream bug, I think there should be an edit in the [run_summarization.py](https://github.com/huggingface/transformers/blob/master/examples/tensorflow/summarization/run_summarization.py) stating its incompatibility with TPU for the timebeing. 

PS: Since, I have not ran the original script, I would like to know whether my above kaggle kernel is missing anything. I was able to run it on GPU. Only got the problem while using TPU.

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.

cc @Rocketknight1 

Hi, I'm sorry for the slow response here! It does seem like an upstream bug, but we'll hopefully be supporting TF 2.6 in the next release. I'm also working on a refactor of the examples using a new data pipeline, so I'll test TPU training with this example when that's implemented to make sure it's working then.

> Hi, I'm sorry for the slow response here! It does seem like an upstream bug, but we'll hopefully be supporting TF 2.6 in the next release. I'm also working on a refactor of the examples using a new data pipeline, so I'll test TPU training with this example when that's implemented to make sure it's working then.

@Rocketknight1 Ohh alright. I will keep this issue open for now since it is not yet solved just incase someone needs it. Eagerly waiting for increased TensorFlow support. :smiley:

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.