Pallas like kernel language for TensorFlow
### Issue type

Feature Request

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

Yes

### Source

source

### TensorFlow version

tf 2.15

### Custom code

Yes

### OS platform and distribution

_No response_

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

JAX provided [Pallas](https://jax.readthedocs.io/en/latest/pallas/tpu.html) to write TPU and GPU kernel. It would be great if TensorFlow had a similar feature.

