Training Vanilla Transformer on TPU gives InternalError
Hi @h4ck4l1 ,

For TPU strategy there is one sharp edge mentioned below.

> To run TF2 programs on TPUs, you can either use .compile and .fit APIs in [tf.keras](https://www.tensorflow.org/api_docs/python/tf/keras) with TPUStrategy, or write your own customized training loop by calling strategy.run directly. Note that TPUStrategy doesn't support pure eager execution, so please make sure the function passed into strategy.run is a [tf.function](https://www.tensorflow.org/api_docs/python/tf/function) or strategy.run is called inside a [tf.function](https://www.tensorflow.org/api_docs/python/tf/function) if eager behavior is enabled.

From the error log in your attached gist the error is generated from [capture_container.capture_by_value()](https://github.com/tensorflow/tensorflow/blob/ab6bbb417f5e336f5f0f4164ade5ac0c7555874c/tensorflow/core/function/capture/capture_container.py#L88).This function will be called  to capture `tensor` if it's external to the graph.

With the above note can you try using `tf.function` decorator for the functions that are defined outside strategy.run and let us know the outcome.

Also I am interested to know why GPU logical device has been created in the failing code.

Thanks! Thanks for responding @SuryanarayanaY !.
I will try your suggestion and will get back to you. Can you show me where GPU logical device was created? I can't find it in the error message, I wish to rectify any redundancies in my code if  any present. 
Q)There are some functions that I am using right now before the strategy.scope() declaration, like to get datasets, to get model...should i decorate everything with @tf.function? or only some specific funcitons?.
Also my extra naive question
Q)Do I require the use the authentication for me to use TPU's?  like we do the auth.authenticate_user() right?
 
I tried your suggestion but i got a new error 
```
---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
[<ipython-input-16-5bd16261e9ac>](https://localhost:8080/#) in <cell line: 1>()
----> 1 model.fit(train_ds,validation_data=valid_ds,epochs=10)

1 frames
[/usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py](https://localhost:8080/#) in error_handler(*args, **kwargs)
     68             # To get the full stack trace, call:
     69             # `tf.debugging.disable_traceback_filtering()`
---> 70             raise e.with_traceback(filtered_tb) from None
     71         finally:
     72             del filtered_tb

[/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py](https://localhost:8080/#) in _numpy(self)
   1107       return self._numpy_internal()
   1108     except core._NotOkStatusException as e:  # pylint: disable=protected-access
-> 1109       raise core._status_to_exception(e) from None  # pylint: disable=protected-access
   1110 
   1111   @property

InvalidArgumentError: Unable to parse tensor proto
```

Please do take a look at the [notebook](https://colab.research.google.com/drive/1lUYapbE1XfLVy7Sc7Xy5oBi-CeBaTRn4#scrollTo=HuUDtZTpOo1Q) maybe it might give an idea?.


















 I would like to add additional data. I am not using gcs for anything here, like in past I had done 
[Notebook with just tensors as input](https://colab.research.google.com/drive/1ii0D2GokcdgxZVjB7oHKy0Lt-I8Ykuin)
[Notebook with tfds datasets as input](https://colab.research.google.com/drive/1DJU058LhhyCfNsuHZ74kZ0E2ziHw-VYo)
both use tpu to train and i got no problem in training them, but this time its slight different its neither tensors nor tfds datasets its custom data pipeline, maybe that is problem(?) but does that mean we can't train custom pipelines on tpu(?) Nevermind after i converted it into tensors it started working again 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/61909">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/61909">No</a>

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

source

### TensorFlow version

2.13.0

### Custom code

Yes

### OS platform and distribution

Linux Ubuntu 22.04.2 (Google Colab)

### Mobile device

Colab

### Python version

3.10.12

### Bazel version

Colab

### GCC/compiler version

Colab

### CUDA/cuDNN version

Colab

### GPU model and memory

Colab

### Current behavior?

Training Transformer model on TPU gives Internal Error 

I have some idea on the error that there's a incompatible tensor ops thats causing the problem but i can't pinpoint it.

I had already done a bigger model which is using pretrained embeddings and it went off without a hitch i tried to replicate the same but with different tfds dataset

If this is already solved please direct me to the relevant links

### Standalone code to reproduce the issue

```
[This is the notebook](https://colab.research.google.com/drive/1y3VEuaYXnsoB42TaVd8UBB-18U8UHFBt#scrollTo=Y0hKZ9yRC3FU)

[This notebook worked fine](https://colab.research.google.com/drive/1DJU058LhhyCfNsuHZ74kZ0E2ziHw-VYo)

Thankyou in advance. I will respond asap
```


### Relevant log output

```shell
---------------------------------------------------------------------------
InternalError                             Traceback (most recent call last)
[<ipython-input-12-013fa12d9e3a>](https://localhost:8080/#) in <cell line: 1>()
----> 1 model.fit(
      2     train_ds,
      3     validation_data=valid_ds,
      4     epochs=EPOCHS,
      5     steps_per_epoch=train_steps,

1 frames
[/usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py](https://localhost:8080/#) in error_handler(*args, **kwargs)
     68             # To get the full stack trace, call:
     69             # `tf.debugging.disable_traceback_filtering()`
---> 70             raise e.with_traceback(filtered_tb) from None
     71         finally:
     72             del filtered_tb

[/usr/local/lib/python3.10/dist-packages/tensorflow/core/function/capture/capture_container.py](https://localhost:8080/#) in capture_by_value(self, graph, tensor, name)
    120         graph_const = self.by_val_internal.get(id(tensor))
    121         if graph_const is None:
--> 122           graph_const = tensor._capture_as_const(name)  # pylint: disable=protected-access
    123           if graph_const is None:
    124             # Some eager tensors, e.g. parallel tensors, are not convertible to

InternalError: failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:35437: Failed to connect to remote host: Connection refused
Additional GRPC error information from remote target /job:localhost/replica:0/task:0/device:CPU:0:
:UNKNOWN:failed to connect to all addresses; last error: UNKNOWN: ipv4:127.0.0.1:35437: Failed to connect to remote host: Connection refused {created_time:"2023-09-19T08:56:09.694479753+00:00", grpc_status:14}
Executing non-communication op <MultiDeviceIteratorInit> originally returned UnavailableError, and was replaced by InternalError to avoid invoking TF network error handling logic.
```
