to_tf_dataset fails on TPU
MEMBER
This might be related to https://github.com/tensorflow/tensorflow/issues/38762 , what do you think @Rocketknight1 ?
> Dataset.from_generator is expected to not work with TPUs as it uses py_function underneath which is incompatible with Cloud TPU 2VM setup. If you would like to read from large datasets, maybe try to materialize it on disk and use TFRecordDataest instead.
2021-11-30T13:32:08Z
https://github.com/huggingface/datasets/issues/3339#issuecomment-982640124
https://api.github.com/repos/huggingface/datasets/issues/3339
IC_kwDODunzps46keX8
https://api.github.com/repos/huggingface/datasets/issues/comments/982640124/reactions
2021-11-30T13:32:08Z
https://api.github.com/repos/huggingface/datasets/issues/comments/982640124
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
MEMBER
Hi @lhoestq @nbroad1881, I think it's very similar, yes. Unfortunately `to_tf_dataset` uses `tf.numpy_function` which can't be compiled - this is a necessary evil to load from the underlying Arrow dataset. We need to update the notebooks/examples to clarify that this won't work, or to identify a workaround. You may be able to get it to work on an actual cloud TPU VM, but those are quite new and we haven't tested it yet. 
2021-12-01T20:04:28Z
https://github.com/huggingface/datasets/issues/3339#issuecomment-984010825
https://api.github.com/repos/huggingface/datasets/issues/3339
IC_kwDODunzps46ptBJ
https://api.github.com/repos/huggingface/datasets/issues/comments/984010825/reactions
2021-12-01T20:04:28Z
https://api.github.com/repos/huggingface/datasets/issues/comments/984010825
https://avatars.githubusercontent.com/u/12866554?v=4
https://api.github.com/users/Rocketknight1/events{/privacy}
https://api.github.com/users/Rocketknight1/followers
https://api.github.com/users/Rocketknight1/following{/other_user}
https://api.github.com/users/Rocketknight1/gists{/gist_id}
https://github.com/Rocketknight1
Rocketknight1
MDQ6VXNlcjEyODY2NTU0
https://api.github.com/users/Rocketknight1/orgs
https://api.github.com/users/Rocketknight1/received_events
https://api.github.com/users/Rocketknight1/repos
https://api.github.com/users/Rocketknight1/starred{/owner}{/repo}
https://api.github.com/users/Rocketknight1/subscriptions
User
https://api.github.com/users/Rocketknight1
public
NONE
Thank you for the explanation. I didn't realize the nuances of `tf.numpy_function`. In this scenario, would it be better to use `export(format='tfrecord')` ?  It's not quite the same, but for very large datasets that don't fit in memory it looks like it is the only option. I haven't used `export` before, but I do recall reading that there are suggestions for how big and how many tfrecords there should be to not bottleneck the TPU. It might be nice if there were a way for the `export` method to split the files up into appropriate chunk sizes depending on the size of the dataset and the number of devices. And if that is too much, it would be nice to be able to specify the number of files that would be created when using `export`. Well... maybe the user should just do the chunking themselves and call `export` a bunch of times.  Whatever the case, you have been helpful. Thanks Tensorflow boy ;-) 
2021-12-01T21:18:07Z
https://github.com/huggingface/datasets/issues/3339#issuecomment-984063079
https://api.github.com/repos/huggingface/datasets/issues/3339
IC_kwDODunzps46p5xn
https://api.github.com/repos/huggingface/datasets/issues/comments/984063079/reactions
2021-12-01T21:18:07Z
https://api.github.com/repos/huggingface/datasets/issues/comments/984063079
https://avatars.githubusercontent.com/u/24982805?v=4
https://api.github.com/users/nbroad1881/events{/privacy}
https://api.github.com/users/nbroad1881/followers
https://api.github.com/users/nbroad1881/following{/other_user}
https://api.github.com/users/nbroad1881/gists{/gist_id}
https://github.com/nbroad1881
nbroad1881
MDQ6VXNlcjI0OTgyODA1
https://api.github.com/users/nbroad1881/orgs
https://api.github.com/users/nbroad1881/received_events
https://api.github.com/users/nbroad1881/repos
https://api.github.com/users/nbroad1881/starred{/owner}{/repo}
https://api.github.com/users/nbroad1881/subscriptions
User
https://api.github.com/users/nbroad1881
public
MEMBER
Yeah, this is something we really should have a proper guide on. I'll make a note to test some things and make a 'TF TPU best practices' notebook at some point, but in the meantime I think your solution of exporting TFRecords will probably work. 
2021-12-02T13:37:36Z
https://github.com/huggingface/datasets/issues/3339#issuecomment-984634984
https://api.github.com/repos/huggingface/datasets/issues/3339
IC_kwDODunzps46sFZo
https://api.github.com/repos/huggingface/datasets/issues/comments/984634984/reactions
2021-12-02T13:37:36Z
https://api.github.com/repos/huggingface/datasets/issues/comments/984634984
https://avatars.githubusercontent.com/u/12866554?v=4
https://api.github.com/users/Rocketknight1/events{/privacy}
https://api.github.com/users/Rocketknight1/followers
https://api.github.com/users/Rocketknight1/following{/other_user}
https://api.github.com/users/Rocketknight1/gists{/gist_id}
https://github.com/Rocketknight1
Rocketknight1
MDQ6VXNlcjEyODY2NTU0
https://api.github.com/users/Rocketknight1/orgs
https://api.github.com/users/Rocketknight1/received_events
https://api.github.com/users/Rocketknight1/repos
https://api.github.com/users/Rocketknight1/starred{/owner}{/repo}
https://api.github.com/users/Rocketknight1/subscriptions
User
https://api.github.com/users/Rocketknight1
public
MEMBER
Also: I knew that tweet would haunt me
2021-12-02T14:21:27Z
https://github.com/huggingface/datasets/issues/3339#issuecomment-984671433
https://api.github.com/repos/huggingface/datasets/issues/3339
IC_kwDODunzps46sOTJ
https://api.github.com/repos/huggingface/datasets/issues/comments/984671433/reactions
2021-12-02T14:21:27Z
https://api.github.com/repos/huggingface/datasets/issues/comments/984671433
https://avatars.githubusercontent.com/u/12866554?v=4
https://api.github.com/users/Rocketknight1/events{/privacy}
https://api.github.com/users/Rocketknight1/followers
https://api.github.com/users/Rocketknight1/following{/other_user}
https://api.github.com/users/Rocketknight1/gists{/gist_id}
https://github.com/Rocketknight1
Rocketknight1
MDQ6VXNlcjEyODY2NTU0
https://api.github.com/users/Rocketknight1/orgs
https://api.github.com/users/Rocketknight1/received_events
https://api.github.com/users/Rocketknight1/repos
https://api.github.com/users/Rocketknight1/starred{/owner}{/repo}
https://api.github.com/users/Rocketknight1/subscriptions
User
https://api.github.com/users/Rocketknight1
public
Using `to_tf_dataset` to create a dataset and then putting it in `model.fit` results in an internal error on TPUs. I've only tried on Colab and Kaggle TPUs, not GCP TPUs.


## Steps to reproduce the bug
I made a colab to show the error. https://colab.research.google.com/drive/12x_PFKzGouFxqD4OuWfnycW_1TaT276z?usp=sharing

## Expected results
dataset from `to_tf_dataset` works in `model.fit`  
Right below the first error in the colab I use `tf.data.Dataset.from_tensor_slices` and `model.fit` works just fine. This is the desired outcome.

## Actual results
```
InternalError: 5 root error(s) found.
  (0) INTERNAL: {{function_node __inference_train_function_30558}} failed to connect to all addresses
Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
:{"created":"@1638231897.932218653","description":"Failed to pick subchannel","file":"third_party/grpc/src/core/ext/filters/client_channel/client_channel.cc","file_line":3151,"referenced_errors":[{"created":"@1638231897.932216754","description":"failed to connect to all addresses","file":"third_party/grpc/src/core/lib/transport/error_utils.cc","file_line":161,"grpc_status":14}]}
	 [[{{node StatefulPartitionedCall}}]]
	 [[MultiDeviceIteratorGetNextFromShard]]
Executing non-communication op <MultiDeviceIteratorGetNextFromShard> originally returned UnavailableError, and was replaced by InternalError to avoid invoking TF network error handling logic.
	 [[RemoteCall]]
	 [[IteratorGetNextAsOptional]]
	 [[tpu_compile_succeeded_assert/_14023832043698465348/_7/_439]]
```

## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version: 1.16.1
- Platform: Linux-5.4.104+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.12
- PyArrow version: 3.0.0
- Tensorflow 2.7.0
- `transformers` 4.12.5
