LSTM loss does not work TPU with bfloat16
@sronen71,
As you mentioned, it was the issue on tensorflow v2.12. When I tried to execute the same mentioned code on tensorflow v2.13`(tpu)` it was executed without any issue/error. Kindly find the gist of it [here](https://colab.research.google.com/gist/tilakrayal/349987540ae6320ed3726d2e86a996d4/untitled1255.ipynb). Thank you! That's great, thank you @tilakrayal !
 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/61352">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/61352">No</a>

### Issue type

Bug

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

Yes

### Source

source

### TensorFlow version

2.12.0

### Custom code

Yes

### OS platform and distribution

Google Colab  + TPU

### 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?

LSTM on TPU works in float32.
Gives error message in bfloat16:
Value passed to parameter 'input' has DataType bfloat16 not in list of allowed values: float16, float32, float64

### Standalone code to reproduce the issue

```shell
Colab Code in this gist:
https://colab.research.google.com/gist/sronen71/cacdc527ea3a7d267e5f47e6dcc8f17f/working_with_rnns.ipynb.

Run with TPU.
```


### Relevant log output

```shell
TypeError                                 Traceback (most recent call last)

<ipython-input-12-d2060b3c689a> in <cell line: 1>()
      1 with strategy.scope():
----> 2   model = build_model()
      3 
      4   model.compile(
      5     loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),

3 frames

/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/op_def_library.py in _SatisfiesTypeConstraint(dtype, attr_def, param_name)
     54     allowed_values = ", ".join(dtypes.as_dtype(x).name for x in allowed_list)
     55     if dtype not in allowed_list:
---> 56       raise TypeError(
     57           f"Value passed to parameter '{param_name}' has DataType "
     58           f"{dtypes.as_dtype(dtype).name} not in list of allowed values: "

TypeError: Exception encountered when calling layer "lstm_4" (type LSTM).

Value passed to parameter 'input' has DataType bfloat16 not in list of allowed values: float16, float32, float64

Call arguments received by layer "lstm_4" (type LSTM):
  • inputs=tf.Tensor(shape=(None, None, 128), dtype=bfloat16)
  • mask=None
  • training=None
  • initial_state=None
```
