Repository: tensorflow
Issue: TF strings do not work on the GPU as indices for  `tf.gather` / `tf.nn.embedding_lookup`
Labels: stat:awaiting tensorflower,type:bug,comp:ops,TF 2.15

### Issue type

Documentation Feature Request

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

No

### Source

binary

### TensorFlow version

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?

As the [`tf.gather`](https://www.tensorflow.org/api_docs/python/tf/gather) documentation suggests, the following code will indeed break on the CPU, but work on the GPU

![image](https://github.com/user-attachments/assets/4c3180fb-92ba-4d34-9518-8af5968f490e)

```py
indices= tf.constant([ 31., 117., 180., 255., 127.,  14.], dtype=tf.float32)

print(indices.shape)

labels = tf.constant([[1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1]], dtype=tf.int32)

print(labels.shape)

tf.nn.embedding_lookup(indices, labels)
```

But this does not to seem to be the case when indices are of type `tf.string` instead:

```py
indices= tf.constant(["a", "b", "c", "d", "e", "f"], dtype=tf.string)

print(indices.shape)

labels = tf.constant([[1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1]], dtype=tf.int32)

print(labels.shape)

tf.nn.embedding_lookup(indices, labels)
```

As the op indeed runs on the CPU (see log output)

A better approach is to indeed use a [`StaticHashTable`](https://www.tensorflow.org/api_docs/python/tf/lookup/StaticHashTable), and then lookup the attribute labels:

```py
text_values= tf.constant(["a", "b", "c", "d", "e", "f"], dtype=tf.string)
text_indices = tf.range(tf.size(text_values), dtype=tf.int32)

text_table = tf.lookup.StaticHashTable(
    initializer = tf.lookup.KeyValueTensorInitializer(
        keys=text_indices,
        values=text_values,
        key_dtype=tf.int32,
        value_dtype=tf.string
    ),
    default_value="",
)

labels = tf.constant([[1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1]], dtype=tf.int32)

text_table.lookup(labels)
```

But I think this should be documented somewhere in `tf.gather` / `tf.nn.embedding_lookup`


Here is a notebook that demonstrates this behaviour:

https://colab.research.google.com/drive/15Ig6Sw39lXRkZ40TvZhs4Aq_v96t-caL?usp=sharing


### Standalone code to reproduce the issue

```shell
indices= tf.constant(["a", "b", "c", "d", "e", "f"], dtype=tf.string)

print(indices.shape)

labels = tf.constant([[1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
        -1, -1, -1, -1]], dtype=tf.int32)

print(labels.shape)

tf.nn.embedding_lookup(indices, labels)
```


### Relevant log output

```shell
(6,)
(1, 100)
---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-9-e96bf33d6da0> in <cell line: 15>()
     13 print(labels.shape)
     14 
---> 15 tf.nn.embedding_lookup(indices, labels)

1 frames
/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py in raise_from_not_ok_status(e, name)
   5881 def raise_from_not_ok_status(e, name) -> NoReturn:
   5882   e.message += (" name: " + str(name if name is not None else ""))
-> 5883   raise core._status_to_exception(e) from None  # pylint: disable=protected-access
   5884 
   5885 

InvalidArgumentError: {{function_node __wrapped__GatherV2_device_/job:localhost/replica:0/task:0/device:CPU:0}} indices[0,1] = -1 is not in [0, 6) [Op:GatherV2] name:
```