InvalidArgumentError: Unsupported data type for TPU: string
@sachinprasadhs,
I tried to execute the mentioned code on tensoflow v2.9 and nightly & it was failing due to `No OpKernel was registered` error. Kindly find the gist of it [here](https://colab.research.google.com/gist/tilakrayal/22ac50ef65e150184ddce505c2bccaba/invalidargumenterror_unsupported_data_type_for_tpu_string.ipynb). @tilakrayal 
That is because, you're trying to initiate tpu-vm from colab. Colab only provides tpu-node, which initialization process is bit different. And as it's mentioned, I've used tpu-vm (on kaggle) and not tpu-node. Node requires to have data on GCP bucket or something and doesn't work on local file system but tpu-vm does. 

However, the actual issue can be reproduced in any system. cc @mattdangerw any thoughts? From TPU [doc](https://cloud.google.com/tpu/docs/troubleshooting/trouble-tf)

> Currently, only the tf.float32, tf.int32, tf.bfloat16, and tf.bool data types are supported on the TPU. Other common data types, such as tf.uint8, tf.string, and tf.int64, must be converted to one of the supported data types during data pre-processing (that is, in the tf.data.Dataset pipeline).

Maybe excluding tokenizer from model (`tfst_model`) and add that to data loader with uniform padding, so that TPU can digest! 

---

**update**

Solved with above approach. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60092">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60092">No</a>

**Environment**

```
OS: Kaggle/Colab
TensorFlow: 2.9.1
```

**Background**

I've a training dataset (`x`: image, `y`: caption). For training, I need to transform this caption to text embedding. And I want to do that inside the model, at training time. And not from data-loader (e.x `tf.data` API). 

Now, including the text-encoder for text embedding inside the actual model (image-captioning model), the whole model can be super heavy. To remedy this, I've followed the following approach. I've tried to transform **text to text embedding** inside the model by overriding the `train_step` method. The process works on GPUs. But for TPU devices, no.

```python
class TextToEmbedding(keras.Model):
    def call(self, inputs):
        return self.model(inputs)
    
    def train_step(self, data):
        x, y = data
        y = tfst_model(y)
        return super().train_step((x, y))
    
    def test_step(self, data):
        x, y = data
        y = tfst_model(y)
        return super().test_step((x, y))
```

On TPUs, the error is showed up. It gives: 

> InvalidArgumentError: Unsupported data type for TPU: string, caused by output Pad_1:0 [Op:__inference_train_function_67882]

Though it's stated that it is unsupported for TPU, but I've failed to understand how that can be unspported on such usual cases. Does string input not support on TPU? 


**Reproducible Code**

[Gist.](https://colab.research.google.com/drive/1C0ejXFVav0yCCDmAATisIhiTFPjGxHKM?usp=sharing)

(Also note, I'm running the code on TPU-vm and not on node).

**Similar issues**

- https://github.com/google/uncertainty-baselines/issues/579
- https://github.com/tensorflow/quantum/issues/707
