Repository: tensorflow
Issue: Gradients can't be computed for keras embeddings
Labels: stat:awaiting tensorflower,type:bug,comp:ops,TF 2.15

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

source

### TensorFlow version

2.15.1

### Custom code

Yes

### OS platform and distribution

Windows 11, Ubuntu 22.04LTS

### Mobile device

_No response_

### Python version

3.11.6

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

We have a problem in shap for quite some while now with embeddings (see [here](https://github.com/shap/shap/issues/3440)). Since we manipulate the graph to adjust the gradient calculation in order to produce shap values we need the layers to be backpropagatable. This does not seem the case for `tensorflow.keras.layers.Embedding` and we do not know a way around this.

(In a previous version there was the possibility to [manipulate](https://github.com/shap/shap/blob/master/shap/explainers/_deep/deep_tf.py#L412-L416) the [`_IsBackpropagatable` function](https://github.com/tensorflow/tensorflow/blob/v1.10.0/tensorflow/python/ops/gradients_impl.py#L293) but this is no longer possible)

In the example below one can see that the gradients just become `None` if the model contains an embedding layer. 
Is there a way around this, so that we can calculate gradients for embeddings again?

### Standalone code to reproduce the issue

```python
import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.datasets import imdb
from tensorflow.keras.preprocessing.sequence import pad_sequences

# Load the IMDb dataset
max_features = 10000  # Only consider the top 10,000 words
maxlen = 100  # Only consider the first 100 words of each movie review

(X_train, y_train), (X_test, y_test) = imdb.load_data(num_words=max_features)
X_train = pad_sequences(X_train, maxlen=maxlen)
X_test = pad_sequences(X_test, maxlen=maxlen)

# Build the model
model = models.Sequential()
embedding_layer = layers.Embedding(input_dim=max_features, output_dim=128, input_length=maxlen)
model.add(embedding_layer)
flat_layer = layers.Flatten()
model.add(flat_layer)
dense_layer = layers.Dense(1, activation='sigmoid')
model.add(dense_layer)

# Build the same model except for the embedding layer
new_model = models.Sequential()
new_model.add(flat_layer)
new_model.add(layers.Dense(1, activation="sigmoid"))

# Forward pass and gradient extraction
@tf.function
def get_gradients_model(inputs):
    inputs = tf.cast(inputs, tf.float32)  # Convert inputs to float32
    with tf.GradientTape() as tape:
        tape.watch(inputs)  # Watch the input tensor to compute gradients w.r.t. it
        predictions = model(inputs)
    
    gradients = tape.gradient(predictions, inputs)
    
    return predictions, gradients

@tf.function
def get_gradients_new_model(inputs):
    inputs = tf.cast(inputs, tf.float32)  # Convert inputs to float32
    with tf.GradientTape() as tape:
        tape.watch(inputs)  # Watch the input tensor to compute gradients w.r.t. it
        predictions = new_model(inputs)
    
    gradients = tape.gradient(predictions, inputs)
    
    return predictions, gradients

# Example usage
sample_input = X_train[:1]  # Select a sample from the training set
sample_label = y_train[:1]  # Corresponding label

predictions, gradients = get_gradients_model(sample_input)
predictions2, gradients2 = get_gradients_new_model(sample_input)
print("Gradients for model:", gradients)
print("Gradients for new_model:", gradients2)
```


### Relevant log output

Gradients for model: None
Gradients for new_model: tf.Tensor(
[[0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
  0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
  0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
  0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
  0. 0. 0. 0.]], shape=(1, 100), dtype=float32)