tf.math.digamma fails when receiving a bfloat16 input
Hi, @mazeltovlee 

I was able to replicate the issue with latest stable `Tensorflow version 2.11` and also with `tf-nightly-2.13.0-dev20230208` and I'm also getting same error which you mentioned in above error log output and for your reference I have added [gist file-tf-2.11 ](https://colab.research.google.com/gist/gaikwadrahul8/39c2058136176a226b5330ff5b04dcd4/-59624-tf-2-11.ipynb)and [gist-file-tf-nightly](https://colab.research.google.com/gist/gaikwadrahul8/aae6b8b60cc0d209ed3daaadc40d74b2/-59624-tf-nightly.ipynb)

As per the `tf.math.digamma()` documentation below it should accept `bfloat16` so it seems like we have to dig more into this issue and we'll update soon here and thank you for noticing the issue with `bfloat16` datatype. Thank you!

```
Computes Psi, the derivative of Lgamma (the log of the absolute value of

Gamma(x)), element-wise.

Args:
  x: A Tensor. Must be one of the following types: bfloat16, half, float32, float64.
  name: A name for the operation (optional).

Returns:
  A Tensor. Has the same type as x.
``` Thanks for your investigation, please let me know if there is something I can do. Hi, @mazeltovlee 

I have executed the same code on `TPU` with `bfloat16` datatype and it seems like it's working as expected and you'll have to use `bfloat16` datatype with `TPU` and it should work so for your reference I have added [gist file ](https://colab.research.google.com/gist/gaikwadrahul8/3df17e2e894684995cd6c57532c38689/-59624.ipynb) and you can try on Google Colab with TPU runtime type.

If issue still persists please let us know or Could you please confirm if this issue is resolved for you ? Please feel free to close the issue if it is resolved ? Thank you! Hi, @gaikwadrahul8 
Thanks for your reply. Since `bfloat16` works fine for the TPU computation and `bfloat16` is not expected to run on other platforms. I will close this issue. 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/59624">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/59624">No</a>
 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/59624">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/59624">No</a>

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

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

Yes

### Source

source

### Tensorflow Version

2.11.0

### Custom Code

Yes

### OS Platform and Distribution

_No response_

### Mobile device

_No response_

### Python version

3.8

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

```shell
This issue is similar to my previous one: https://github.com/tensorflow/tensorflow/issues/59611.
Following the documentation: https://www.tensorflow.org/api_docs/python/tf/math/digamma, tf.math.digamma should accept the bfloat16 input but it fails. BTW the sentence in the documentation online "Computes Psi, the derivative of Lgamma (the log of the absolute value of" seems incomplete
```


### Standalone code to reproduce the issue

```shell
import tensorflow as tf
tf.math.digamma(tf.Variable(0., dtype="bfloat16"))
```


### Relevant log output

```shell
tensorflow.python.framework.errors_impl.NotFoundError: Could not find device for node: {{node Digamma}} = Digamma[T=DT_BFLOAT16]
All kernels registered for op Digamma:
  device='XLA_CPU_JIT'; T in [DT_FLOAT, DT_DOUBLE, DT_BFLOAT16, DT_HALF]
  device='XLA_GPU_JIT'; T in [DT_FLOAT, DT_DOUBLE, DT_BFLOAT16, DT_HALF]
  device='GPU'; T in [DT_DOUBLE]
  device='GPU'; T in [DT_FLOAT]
  device='GPU'; T in [DT_HALF]
  device='CPU'; T in [DT_DOUBLE]
  device='CPU'; T in [DT_HALF]
  device='CPU'; T in [DT_FLOAT]
 [Op:Digamma]
```
</details>