InvalidArgumentError: Dst node should be assigned to an allowed device. [TPU]
Any update? This program works fine on single and mult-gpu devices. But it raises issue for running on TPU node (perhaps due to unknown shape initiated in init!!).  This program exhibits two behaviors presently unsupported by TPUs:

1) Variables with unknown shapes.
   a) "TPUs do not have registered OpKernel support for [tf.raw_ops.]Shape."
2) Variables whose shapes change over time.

This error is appearing due to the first behavior, but resolving that will induce the second.

My guidance would be to rewrite this program to avoid these behaviors, if possible. For example, you could place these variables on the host CPU. Set it true but still gives error.

```
tf.config.set_soft_device_placement(
    enabled
)
```

How to place on CPU devices (colab or kaggle)? 

```
tf.config.list_logical_devices('TPU')
```

Only returns the available tpu node. I also tried `tf.device('CPU'):`, but no luck.  The simplest way to place them on the host device is to move the `tf.Variable` statements outside of the distribution strategy scope. That would require a little bit of rearranging here.

This assumes that you cannot make the variables fixed shape and maintain the TPU placement. There are various ways to achieve that, like pre-allocating them to be the largest possible size (after all the `tf.concat` ops) or using another aggregation method (e.g., the difference between concat-sum-divide and moving averages). You may be able to use a `tf.TensorArray` if you know an upper-bound for the number of concatenations; however, I'm unsure whether that would apply here. > The simplest way to place them on the host device is to move the tf.Variable statements outside of the distribution strategy scope. That would require a little bit of rearranging here.

Even though I place variable outside the scope, it complains the same. Any catch?

https://colab.research.google.com/drive/18O9M6kiD-1kyrodS5_PsCshgW08xsKst?usp=sharing Ah, I see. There's also some trickiness here because these are distributed/per-device values being sent to a single device. I'll dig into this further and see what I can turn up. A workaround might involve writing a [custom training loop](https://www.tensorflow.org/tutorials/distribute/custom_training) that reduces the per-device values ([link](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy#reduce)), and stores them host-side.

A broader question--are per-batch, per-device values necessary? If they can be averaged or, otherwise, aggregated that would save a lot of memory and fit into the existing patterns (Keras metrics). > A broader question--are per-batch, per-device values necessary? If they can be averaged or, otherwise, aggregated that would save a lot of memory and fit into the existing patterns (Keras metrics).

My goal here is to get the `y`, `y_pred` of validation data set inside the callback API (`on_epoch_end`) to perform some post-processing stuff. Originally posted [here](https://github.com/keras-team/tf-keras/issues/330). It is [suggested](https://github.com/keras-team/tf-keras/issues/330) to use the way I'm doing (above colab). So, yes, per-batch, per-device data is necessary.  Apologies for the delayed response.

One potential solution may be to use an all-gather to collect the per-device values and then write a summary file via outside compilation. Writing a summary file could consume less memory during training, but comes at the cost of having to load the data back later. Here is a loose example of what that could look like: [link](https://gist.github.com/jszaday/ca0d3eb05c487bd3e2a74b697b22c3f0).

A couple things to note:
- TPU workers cannot access a Colab's local storage, only GCS buckets (in case of `File system scheme '[local]' not implemented...`).
- This example does not enable "soft device placement" since it uses the TF1 equivalent (outside compilation).
- One could increment `step` in a Keras callback like `on_train_batch_begin`.

I hope this example illustrates a potential way forward for your issue; at the very least, its pattern could be adapted to store the per-batch, per-device values another way. @jszaday Thanks for the update with details explanations. It looks interesting, let me check and get back to you. This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. @jszaday 
Sorry for the delay. Though the provided gist looks a great starting point, I would like to have more simplified APIs for this. Changing the end user code for the target results (batch/epoch level output) on TPU looks heavy though it should be run as it is without any hacks. I wonder how pytorch-lightening (pl) does this things in back. In pl, the callback APIs does output data step/epoch wise and pl code can be run on gpu/tpu devices with ease. 

Recently this FR is reoccurred in this PR https://github.com/keras-team/keras-io/pull/1314

```python
class EvaluateCOCOMetricsCallback(keras.callbacks.Callback):
    def __init__(self, data):
        super().__init__()
        self.data = data
        self.metrics = keras_cv.metrics.BoxCOCOMetrics(
            bounding_box_format="xywh",
            # passing 1e9 ensures we never evaluate until
            # `metrics.result(force=True)` is
            # called.
            evaluate_freq=1e9,
        )

    def on_epoch_end(self, epoch, logs):
        self.metrics.reset_state()
        for batch in tqdm.tqdm(self.data):
            images, y_true = batch[0], batch[1]
            y_pred = self.model.predict(images, verbose=0)
            self.metrics.update_state(y_true, y_pred)

        metrics = self.metrics.result(force=True)
        logs.update(metrics)
        return logs

model.fit(
    train_ds.take(20),
    validation_data=eval_ds.take(20),
    # Run for 10-35~ epochs to achieve good scores.
    epochs=1,
    callbacks=[EvaluateCOCOMetricsCallback(eval_ds.take(20))],
)
```

For TPU, it's written as follows, where the same validation data set (`eval_ds` is passed to the model again via callback, even though it's already processed via `validation_data`).

cc. @LukeWood  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/59511">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/59511">No</a>

 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

Yes

### Source

source

### Tensorflow Version

2.4.1

### Custom Code

Yes

### OS Platform and Distribution

Kaggle

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

> InvalidArgumentError: Dst node should be assigned to an allowed device. Found an edge from node concat_variable_140022353296720_handle_inputs_0/shape/_4 assigned to /job:localhost/replica:0/task:0/device:COMPOSITE:0 to node TPUReplicate/_compile/_7103042970398181355/_7 assigned to /job:worker/replica:0/task:0/device:CPU:0

### Standalone code to reproduce the issue

```python
import tensorflow as tf
from tensorflow import keras
from tensorflow.experimental import numpy as tnp
import numpy as np

def set_tpu(mixed_precision=True):
        tpu = tf.distribute.cluster_resolver.TPUClusterResolver.connect() 
        if mixed_precision:
            keras.mixed_precision.set_global_policy("mixed_bfloat16") 
        tf.config.set_soft_device_placement(False)
        strategy = tf.distribute.TPUStrategy(tpu)
        physical_devices = tf.config.list_logical_devices('TPU')
        return (strategy, physical_devices)

    
mxp = False
jit = False
strategy, physical_devices = set_tpu(mixed_precision=mxp)
physical_devices, tf.__version__
```
```python
class CustomModel(keras.Model):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.val_x = tf.Variable((
            tnp.empty((0, 32), dtype=tf.float32)), shape=[None, 32]
        )
        self.val_gt = tf.Variable(
            tnp.empty((0), dtype=tf.float32), shape=[None]
        )
        self.val_pred = tf.Variable(
            tnp.empty((0, 1), dtype=tf.float32), shape=[None, 1]
        )
        
    def test_step(self, data):
        x, y = data
        y_pred = self(x, training=False)
        self.compiled_loss(y, y_pred, regularization_losses=self.losses)
        self.compiled_metrics.update_state(y, y_pred)
        
        # ATTENTION 
        # Main Cause !!!
        self.val_x.assign(
            tf.concat([self.val_x, x], axis=0)
        )
        self.val_gt.assign(
            tf.concat([self.val_gt, y], axis=0)
        )
        self.val_pred.assign(
            tf.concat([self.val_pred, y_pred], axis=0)
        )
        return {m.name: m.result() for m in self.metrics}
```
```python
with strategy.scope():
    inputs = keras.Input(shape=(32,))
    outputs = keras.layers.Dense(1, dtype='float32')(inputs)
    model = CustomModel(inputs, outputs)
    model.compile(
        optimizer="adam", loss="mse", metrics=["mae"],
    )

x = np.random.random((1000, 32))
y = np.random.random((1000))
x_test = np.random.random((10, 32))
y_test = np.random.random((10))

model.fit(
    x, 
    y, 
    epochs=5, 
    validation_data=(x_test, y_test),
    verbose=2, 
    batch_size=4,
)
```

### Relevant log output

> 2023-02-01 11:55:11.789825: W ./tensorflow/core/distributed_runtime/eager/destroy_tensor_handle_node.h:57] Ignoring an error encountered when deleting remote tensors handles: Invalid argument: Unable to find the relevant tensor remote_handle: Op ID: 1206, Output num: 0
Additional GRPC error information from remote target /job:worker/replica:0/task:0:
:{"created":"@1675252511.789496486","description":"Error received from peer ipv4:10.0.0.2:8470","file":"external/com_github_grpc_grpc/src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Unable to find the relevant tensor remote_handle: Op ID: 1206, Output num: 0","grpc_status":3}
Epoch 1/5
2023-02-01 11:55:15.394719: W tensorflow/core/distributed_runtime/eager/remote_tensor_handle_data.cc:76] Unable to destroy remote tensor handles. If you are running a tf.function, it usually indicates some op in the graph gets an error: Dst node should be assigned to an allowed device. Found an edge from node concat_variable_140022353296720_handle_inputs_0/shape/_4 assigned to /job:localhost/replica:0/task:0/device:COMPOSITE:0 to node TPUReplicate/_compile/_7103042970398181355/_7 assigned to /job:worker/replica:0/task:0/device:CPU:0

## Others

- Program works fine in GPU / CPU
- Relevant ticket. https://github.com/keras-team/tf-keras/issues/330