tf.linalg.diag_part() is slow on TPU
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/60166">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/60166">No</a>
 @johnypark 
I tried to reproduce the issue on Colab using TF v2.12 but facing a different error. Could you please refer to the gist [here](https://colab.research.google.com/gist/tiruk007/872b0bc68dcffe20b7351a5eb9eeafdb/untitled177.ipynb) for reference. Could you please provide detailed steps  to replicate the issue reported here.


Thank 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.
<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Performance

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

No

### Source

source

### Tensorflow Version

2.11.0

### Custom Code

Yes

### OS Platform and Distribution

Google colab

### Mobile device

_No response_

### Python version

3.9.16

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?


I am working on implementing efficient attention mechanisms using tensorflow. (such as longformer).

The local interaction in the longformer only requires tokens to interact with their neighboring ones, which could reduce the computation overhead calculating the full dot product between two tensors.

The problem I am currently experiencing is that tf.linalg.diag_part(), a function used for extracting the band diagonal from the key tensor, works extremely slow on TPU. 

In my test, the efficient attention takes 52sec to compute, whereas it takes 2sec to compute the regular dot product attention.

I tried warping it on @tf.function decorator jit_compile=True, but it is still slower than the full dot product calculation.

Using CPU, the efficient attention gets x4-5 performance boost depending on the settings.

There may be some limitations in tf.linalg.diag_part() in parallel use. The source code is written in C++, which I lack expertise in. 

Is there any way to enhance the performance of tf.linalg.diag_part() in parallel processing?


### Standalone code to reproduce the issue

```shell
import tensorflow as tf
from functools import partial

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))
strategy = tf.distribute.TPUStrategy(resolver)


def split_head(input, B, N, D, num_heads):
    assert D/num_heads == D//num_heads, "D must be divisible by num_heads"
   
    x = tf.reshape(input, 
               (B, N, num_heads, D//num_heads))
    return tf.transpose(x, perm = (0, 2, 1, 3))
       
B = 2048
N = 1024
D = 1024
H = 32
b_k = 15
input = tf.ones((B, N, D))

with strategy.scope():
      
 #tf.reshape(tf.range(B*N*D), (B,N,D))
  q = input
  k = input
  split_head = partial(split_head, B= B, N=N, D= D)
  q = split_head(q, num_heads= H)
  k = split_head(k, num_heads= H)
  print(B,N,D)
  k_T = tf.transpose(k, perm =(0, 1, 3, 2))
  print(k_T.shape)
  band_k = extract_band(input = k, b_k = b_k)
  print(band_k.shape)
  if N >= D//H:
      attention = q@tf.transpose(band_k, perm = (0, 1, 3, 2))
      print(attention.shape)
      attention = tf.reduce_mean(attention, axis = 1) # combine head
      attention_map = tf.linalg.diag(tf.transpose(attention, perm = (0, 2, 1) ), k = (-b_k//2, b_k//2))
```


### Relevant log output

_No response_</details>