Repository: tensorflow
Issue: CTC Loss errors on 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

source

### TensorFlow version

2.12

### Custom code

Yes

### OS platform and distribution

Kaggle notebook

### Mobile device

_No response_

### Python version

3.10

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

Keras Model with LSTM+ CTC loss is running normally on GPU,
 but on VM TPU it  prints grappler errors. The errors usually don't stop the execution but the loss is nan.
It sometimes also crashes with coredump, but it's not consistent.

I created a public kaggle notebook with the code producing the issue here:
https://www.kaggle.com/code/shaironen/ctc-example/notebook
The grappler errors are:

E tensorflow/core/grappler/optimizers/meta_optimizer.cc:954] model_pruner failed: INVALID_ARGUMENT: Graph does not contain terminal node AssignAddVariableOp.
2023-07-17 09:06:37.445931: E tensorflow/core/grappler/optimizers/meta_optimizer.cc:954] model_pruner failed: INVALID_ARGUMENT: Graph does not contain terminal node AssignAddVariableOp.
E tensorflow/core/grappler/optimizers/meta_optimizer.cc:954] model_pruner failed: INVALID_ARGUMENT: Graph does not contain terminal node AssignAddVariableOp.

In addition, I read it's recommended to use model.compile(jit_compile=True) while on GPU to pre-diagnose TPU issues. It gives similar errors and terminates.
(with jit_compile=False it run normally on GPU only).

According to  tf.nn.ctc_loss documentation it should work on tpu.


### Standalone code to reproduce the issue

```shell
https://www.kaggle.com/code/shaironen/ctc-example/notebook
```


### Relevant log output

```shell
WARNING:tensorflow:From /usr/local/lib/python3.8/site-packages/tensorflow/python/ops/ctc_ops.py:1512: alias_inplace_add (from tensorflow.python.ops.inplace_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Prefer tf.tensor_scatter_nd_add, which offers the same functionality with well-defined read-write semantics.

WARNING:tensorflow:From /usr/local/lib/python3.8/site-packages/tensorflow/python/ops/ctc_ops.py:1512: alias_inplace_add (from tensorflow.python.ops.inplace_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Prefer tf.tensor_scatter_nd_add, which offers the same functionality with well-defined read-write semantics.

WARNING:tensorflow:From /usr/local/lib/python3.8/site-packages/tensorflow/python/ops/ctc_ops.py:1495: alias_inplace_update (from tensorflow.python.ops.inplace_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Prefer tf.tensor_scatter_nd_update, which offers the same functionality with well-defined read-write semantics.

WARNING:tensorflow:From /usr/local/lib/python3.8/site-packages/tensorflow/python/ops/ctc_ops.py:1495: alias_inplace_update (from tensorflow.python.ops.inplace_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Prefer tf.tensor_scatter_nd_update, which offers the same functionality with well-defined read-write semantics.
2023-07-17 09:29:45.795306: E tensorflow/core/grappler/optimizers/meta_optimizer.cc:954] model_pruner failed: INVALID_ARGUMENT: Graph does not contain terminal node AssignAddVariableOp.
2023-07-17 09:29:46.168913: E tensorflow/core/grappler/optimizers/meta_optimizer.cc:954] model_pruner failed: INVALID_ARGUMENT: Graph does not contain terminal node AssignAddVariableOp.

 99/100 [============================>.] - ETA: 0s - loss: nan

2023-07-17 09:30:01.881305: E tensorflow/core/grappler/optimizers/meta_optimizer.cc:954] model_pruner failed: INVALID_ARGUMENT: Graph does not contain terminal node StatefulPartitionedCall.
2023-07-17 09:30:02.118121: E tensorflow/core/grappler/optimizers/meta_optimizer.cc:954] model_pruner failed: INVALID_ARGUMENT: Graph does not contain terminal node StatefulPartitionedCall.
F0717 09:30:03.134464    3526 throw_delegate.cc:121] RAW: absl::container_internal::raw_hash_map<>::at
    @     0x78bd48329338  (unknown)
    @     0x78bd483271e6  (unknown)
    @     0x78bd465e0f20  (unknown)
    @     0x78bd465d6231  (unknown)
    @     0x78bd465d5b59  (unknown)
    @     0x78bd476ce3c9  (unknown)
    @     0x78bd476cf3d7  (unknown)
    @     0x78bd476ccfba  (unknown)
    @     0x78bd40ef00e6  (unknown)
    @     0x78bd465d4e1f  (unknown)
    @     0x78bd465d6919  (unknown)
    @     0x78bd465d662c  (unknown)
    @     0x78bd46345238  (unknown)
    @     0x78bd430f4143  (unknown)
    @     0x78bd46d5f6b4  (unknown)
    @     0x78bd46d5dc50  (unknown)
    @     0x78bd46d5d54e  (unknown)
    @     0x78bd430cee43  (unknown)
    @     0x78bd430de27a  (unknown)
    @     0x78bd430ee05b  (unknown)
    @     0x78bd430ee945  (unknown)
    @     0x78bd424648a3  (unknown)
    @     0x78bd4245e54c  (unknown)
    @     0x78bd41562909  (unknown)
    @     0x78bd41563a59  (unknown)
    @     0x78bd40f32a1c  TpuCompile_CompileAndBuild
    @     0x78bd515e0225  tensorflow::tpu::TpuProgramGroup::CompileAndBuild()
    @     0x78bd51530d6f  tensorflow::tpu::TpuCompileOpKernelImpl::Compile()
    @     0x78bd515e34df  tensorflow::tpu::TpuCompileOpKernelCommon::CompileLocallyAndFillHostCacheInternal()
    @     0x78bd515e3aa1  tensorflow::tpu::TpuCompileOpKernelCommon::CompileLocallyAndFillHostCache()
    @     0x78bd515e3c8b  tensorflow::tpu::TpuCompileOpKernelCommon::ComputeInternal()::{lambda()#3}::operator()()
    @     0x78bd515e3d6c  std::_Function_handler<>::_M_invoke()
    @     0x78bd515cc070  tensorflow::tpu::TpuCompilationCacheExternal::InitializeEntry()
    @     0x78bd51612e72  tensorflow::tpu::TpuCompilationCacheInterface::CompileIfKeyAbsentHelper()
    @     0x78bd5161391a  tensorflow::tpu::TpuCompilationCacheInterface::CompileIfKeyAbsent()
    @     0x78bd515e4323  tensorflow::tpu::TpuCompileOpKernelCommon::ComputeInternal()
    @     0x78bd515e6fc4  tensorflow::tpu::TpuCompileOpKernelCommon::Compute()
    @     0x78bd62de6b6d  tensorflow::ThreadPoolDevice::Compute()
    @     0x78bd62ecfe2c  tensorflow::(anonymous namespace)::ExecutorState<>::Process()
    @     0x78bd62eb9272  std::_Function_handler<>::_M_invoke()
    @     0x78bd621b7275  Eigen::ThreadPoolTempl<>::WorkerLoop()
    @     0x78bd621b41c7  std::_Function_handler<>::_M_invoke()
    @     0x78bd62cefabf  tsl::(anonymous namespace)::PThread::ThreadFn()
    @     0x78be31abeea7  start_thread
https://symbolize.stripped_domain/r/?trace=78bd48329338,78bd483271e5,78bd465e0f1f,78bd465d6230,78bd465d5b58,78bd476ce3c8,78bd476cf3d6,78bd476ccfb9,78bd40ef00e5,78bd465d4e1e,78bd465d6918,78bd465d662b,78bd46345237,78bd430f4142,78bd46d5f6b3,78bd46d5dc4f,78bd46d5d54d,78bd430cee42,78bd430de279,78bd430ee05a,78bd430ee944,78bd424648a2,78bd4245e54b,78bd41562908,78bd41563a58,78bd40f32a1b,78bd515e0224,78bd51530d6e,78bd515e34de,78bd515e3aa0,78bd515e3c8a,78bd515e3d6b,78bd515cc06f,78bd51612e71,78bd51613919,78bd515e4322,78bd515e6fc3,78bd62de6b6c,78bd62ecfe2b,78bd62eb9271,78bd621b7274,78bd621b41c6,78bd62cefabe,78be31abeea6&map=aef7fe2e538f701f46d88df9ee3b51d79ec62b1e:78bd6181e000-78bd637f9728,8f79f803f683427be94b1cfeea32716e6ef365e4:78bd48eb8000-78bd60e0d830,1278088d049ad36cb636fbbc76303cb3:78bd3cc43000-78bd484907c0 
https://symbolize.stripped_domain/r/?trace=78be31b11ce1,78be31b11d5f,78bd48329337,78bd483271e5,78bd465e0f1f,78bd465d6230,78bd465d5b58,78bd476ce3c8,78bd476cf3d6,78bd476ccfb9,78bd40ef00e5,78bd465d4e1e,78bd465d6918,78bd465d662b,78bd46345237,78bd430f4142,78bd46d5f6b3,78bd46d5dc4f,78bd46d5d54d,78bd430cee42,78bd430de279,78bd430ee05a,78bd430ee944,78bd424648a2,78bd4245e54b,78bd41562908,78bd41563a58,78bd40f32a1b,78bd515e0224,78bd51530d6e,78bd515e34de,78bd515e3aa0,78bd515e3c8a&map=8f79f803f683427be94b1cfeea32716e6ef365e4:78bd48eb8000-78bd60e0d830,1278088d049ad36cb636fbbc76303cb3:78bd3cc43000-78bd484907c0 
*** SIGABRT received by PID 2614 (TID 3526) on cpu 51 from PID 2614; ***
E0717 09:30:03.650583    3526 coredump_hook.cc:414] RAW: Remote crash data gathering hook invoked.
E0717 09:30:03.650603    3526 client.cc:278] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
E0717 09:30:03.650607    3526 coredump_hook.cc:512] RAW: Sending fingerprint to remote end.
E0717 09:30:03.650615    3526 coredump_socket.cc:120] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket
E0717 09:30:03.650619    3526 coredump_hook.cc:518] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?
E0717 09:30:03.650623    3526 coredump_hook.cc:580] RAW: Dumping core locally.
```