Repository: tensorflow
Issue: tf.debugging.experimental.enable_dump_debug_info (Debugger V2) error with TPU
Labels: stat:awaiting tensorflower,type:bug,comp:tpus,TF 2.12

### Issue type

Bug

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

Yes

### Source

binary

### TensorFlow version

tf 2.12 (but also 2.14 nightly)

### Custom code

Yes

### OS platform and distribution

Google Colab

### Mobile device

no

### Python version

3.10

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

N/A (using TPU)

### GPU model and memory

_No response_

### Current behavior?

I'm trying to use the `tf.debugging.experimental.enable_dump_debug_info(...)` function with TPU.

I've tested my code without the TPU strategy bit, and it works, I can use the Debugger V2 fine.

I've also tested the code without the debugger bit and it trains fine as well.

But together it gives me this error:

```
Traceback (most recent call last):
  File "/content/tbscript.py", line 47, in <module>
    model = train()
  File "/content/tbscript.py", line 40, in train
    model.fit(x=x_train, 
  File "/usr/local/lib/python3.10/dist-packages/keras/src/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/debug/lib/dumping_callback.py", line 579, in <listcomp>
    output_tensor_device_ids = [writer.RegisterDeviceAndGetId(output.device)
ValueError: Cannot assign a device for operation IteratorGetNextAsOptional: Could not satisfy explicit device specification '' because the node {{colocation_node IteratorGetNextAsOptional}} was colocated with a group of nodes that required incompatible device '/job:worker/replica:0/task:0/device:TPU:0'. All available devices [/job:worker/replica:0/task:0/device:CPU:0, /job:worker/replica:0/task:0/device:TPU:0, /job:worker/replica:0/task:0/device:TPU:1, /job:worker/replica:0/task:0/device:TPU:2, /job:worker/replica:0/task:0/device:TPU:3, /job:worker/replica:0/task:0/device:TPU:4, /job:worker/replica:0/task:0/device:TPU:5, /job:worker/replica:0/task:0/device:TPU:6, /job:worker/replica:0/task:0/device:TPU:7, /job:worker/replica:0/task:0/device:TPU_SYSTEM:0, /job:worker/replica:0/task:0/device:XLA_CPU:0, /job:localhost/replica:0/task:0/device:CPU:0, /job:localhost/replica:0/task:0/device:COMPOSITE:0]. 
Colocation Debug Info:
Colocation group had the following types and supported devices: 
Root Member(assigned_device_name_index_=2 requested_device_name_='/job:worker/replica:0/task:0/device:TPU:0' assigned_device_name_='/job:worker/replica:0/task:0/device:TPU:0' resource_device_name_='/job:worker/replica:0/task:0/device:TPU:0' supported_device_types_=[CPU] possible_devices_=[]
OptionalGetValue: CPU TPU XLA_CPU 
DebugNumericSummaryV2: CPU 
IteratorGetNext: CPU TPU XLA_CPU 
Identity: CPU TPU TPU_SYSTEM XLA_CPU COMPOSITE 
Switch: CPU TPU TPU_SYSTEM XLA_CPU COMPOSITE 
IteratorGetNextAsOptional: CPU TPU XLA_CPU 
DebugIdentityV2: CPU TPU TPU_SYSTEM XLA_CPU COMPOSITE 
OptionalHasValue: CPU TPU XLA_CPU 
_Arg: CPU TPU TPU_SYSTEM XLA_CPU COMPOSITE 

Colocation members, user-requested devices, and framework assigned devices, if any:
  iterator_1 (_Arg)  framework assigned device=/job:worker/replica:0/task:0/device:TPU:0
  IteratorGetNextAsOptional (IteratorGetNextAsOptional) 
  OptionalHasValue (OptionalHasValue) 
  cond/IteratorGetNextAsOptional/_5 (Switch) 
  cond/iterator_1/_13 (Switch) 
  Func/cond/then/_0/input/_39 (Identity) 
  Func/cond/then/_0/input/_47 (Identity) 
  cond/then/_0/cond/OptionalHasValue (OptionalHasValue) 
  Func/cond/else/_1/input/_73 (Identity) 
  Func/cond/else/_1/input/_81 (Identity) 
  cond/else/_1/cond/IteratorGetNext (IteratorGetNext) 
  cond/else/_1/cond/IteratorGetNext/DebugNumericSummaryV2 (DebugNumericSummaryV2) 
  cond/else/_1/cond/IteratorGetNext/DebugIdentityV2_1511 (DebugIdentityV2) 
  cond/then/_0/cond/cond/Func/cond/then/_0/input/_39/_111 (Switch) 
  Func/cond/then/_0/cond/cond/then/_106/input/_179 (Identity) 
  cond/then/_0/cond/cond/then/_106/cond/cond/OptionalGetValue (OptionalGetValue) /job:worker/replica:0/task:0/device:TPU:0
  cond/then/_0/cond/cond/then/_106/cond/cond/OptionalGetValue/DebugNumericSummaryV2 (DebugNumericSummaryV2) /job:worker/replica:0/task:0/device:TPU:0
  cond/then/_0/cond/cond/then/_106/cond/cond/OptionalGetValue/DebugIdentityV2_1369 (DebugIdentityV2) /job:worker/replica:0/task:0/device:TPU:0
  Func/cond/then/_0/cond/cond/else/_107/input/_184 (Identity) 

	 [[{{node IteratorGetNex ... [truncated]
Exception ignored in atexit callback: <function async_wait at 0x7ecd97fc5d80>
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/eager/context.py", line 2796, in async_wait
    context().sync_executors()
  File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/eager/context.py", line 742, in sync_executors
    pywrap_tfe.TFE_ContextSyncExecutors(self._context_handle)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Cannot assign a device for operation IteratorGetNextAsOptional: Could not satisfy explicit device specification '' because the node {{colocation_node IteratorGetNextAsOptional}} was colocated with a group of nodes that required incompatible device '/job:worker/replica:0/task:0/device:TPU:0'. All available devices [/job:worker/replica:0/task:0/device:CPU:0, /job:worker/replica:0/task:0/device:TPU:0, /job:worker/replica:0/task:0/device:TPU:1, /job:worker/replica:0/task:0/device:TPU:2, /job:worker/replica:0/task:0/device:TPU:3, /job:worker/replica:0/task:0/device:TPU:4, /job:worker/replica:0/task:0/device:TPU:5, /job:worker/replica:0/task:0/device:TPU:6, /job:worker/replica:0/task:0/device:TPU:7, /job:worker/replica:0/task:0/device:TPU_SYSTEM:0, /job:worker/replica:0/task:0/device:XLA_CPU:0, /job:localhost/replica:0/task:0/device:CPU:0, /job:localhost/replica:0/task:0/device:COMPOSITE:0]. 
Colocation Debug Info:
Colocation group had the following types and supported devices: 
Root Member(assigned_device_name_index_=2 requested_device_name_='/job:worker/replica:0/task:0/device:TPU:0' assigned_device_name_='/job:worker/replica:0/task:0/device:TPU:0' resource_device_name_='/job:worker/replica:0/task:0/device:TPU:0' supported_device_types_=[CPU] possible_devices_=[]
OptionalGetValue: CPU TPU XLA_CPU 
DebugNumericSummaryV2: CPU 
IteratorGetNext: CPU TPU XLA_CPU 
Identity: CPU TPU TPU_SYSTEM XLA_CPU COMPOSITE 
Switch: CPU TPU TPU_SYSTEM XLA_CPU COMPOSITE 
IteratorGetNextAsOptional: CPU TPU XLA_CPU 
DebugIdentityV2: CPU TPU TPU_SYSTEM XLA_CPU COMPOSITE 
OptionalHasValue: CPU TPU XLA_CPU 
_Arg: CPU TPU TPU_SYSTEM XLA_CPU COMPOSITE 

Colocation members, user-requested devices, and framework assigned devices, if any:
  iterator_1 (_Arg)  framework assigned device=/job:worker/replica:0/task:0/device:TPU:0
  IteratorGetNextAsOptional (IteratorGetNextAsOptional) 
  OptionalHasValue (OptionalHasValue) 
  cond/IteratorGetNextAsOptional/_5 (Switch) 
  cond/iterator_1/_13 (Switch) 
  Func/cond/then/_0/input/_39 (Identity) 
  Func/cond/then/_0/input/_47 (Identity) 
  cond/then/_0/cond/OptionalHasValue (OptionalHasValue) 
  Func/cond/else/_1/input/_73 (Identity) 
  Func/cond/else/_1/input/_81 (Identity) 
  cond/else/_1/cond/IteratorGetNext (IteratorGetNext) 
  cond/else/_1/cond/IteratorGetNext/DebugNumericSummaryV2 (DebugNumericSummaryV2) 
  cond/else/_1/cond/IteratorGetNext/DebugIdentityV2_1511 (DebugIdentityV2) 
  cond/then/_0/cond/cond/Func/cond/then/_0/input/_39/_111 (Switch) 
  Func/cond/then/_0/cond/cond/then/_106/input/_179 (Identity) 
  cond/then/_0/cond/cond/then/_106/cond/cond/OptionalGetValue (OptionalGetValue) /job:worker/replica:0/task:0/device:TPU:0
  cond/then/_0/cond/cond/then/_106/cond/cond/OptionalGetValue/DebugNumericSummaryV2 (DebugNumericSummaryV2) /job:worker/replica:0/task:0/device:TPU:0
  cond/then/_0/cond/cond/then/_106/cond/cond/OptionalGetValue/DebugIdentityV2_1369 (DebugIdentityV2) /job:worker/replica:0/task:0/device:TPU:0
  Func/cond/then/_0/cond/cond/else/_107/input/_184 (Identity) 

	 [[{{node IteratorGetNex ... [truncated]
2023-07-28 06:05:18.174660: W ./tensorflow/core/distributed_runtime/eager/destroy_tensor_handle_node.h:59] Ignoring an error encountered when deleting remote tensors handles: INVALID_ARGUMENT: Unable to find the relevant tensor remote_handle: Op ID: 900, Output num: 0
Additional GRPC error information from remote target /job:worker/replica:0/task:0 while calling /tensorflow.eager.EagerService/Enqueue:
:{"created":"@1690524318.171303502","description":"Error received from peer ipv4:10.15.76.74: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: 900, Output num: 0","grpc_status":3} [type.googleapis.com/tensorflow.core.platform.ErrorSourceProto='\x08\x05']
```

yeah very long.

oh and I've tried `tf.config.set_soft_device_placement(True)` but no result

### Standalone code to reproduce the issue

```shell
Here's a reproducible test case for getting the error:
https://colab.research.google.com/drive/169agwcqy3-M8hQnSAx5EAr64ya2bSMl3?usp=sharing

but you're not supposed to run it on colab, as I've noticed the `tf.debugging.experimental.enable_dump_debug_info` function generally doesn't work with it. So I put it in a .py script and run it with !python3. Without the TPU part it works fine.
```


### Relevant log output

_No response_