[XLA] TF XLA outputs abnormal value when compiling Embedding #82317
Description
Issue type
Bug
Have you reproduced the bug with TensorFlow Nightly?
Yes
Source
source
TensorFlow version
nightly
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?
For Embedding operator, when I set the input_dim=1 (which means it indexed from 0 to 0), the output always returns 0 without XLA.
After compilation, the outputs are usually some random tensors.
Standalone code to reproduce the issue
import tensorflow as tf
tf.random.set_seed(42)
x = tf.constant([1])
# uncompiled model
class Model(tf.keras.Model):
def __init__(self):
super(Model, self).__init__()
self.embedding = tf.keras.layers.Embedding(1, 1)
def call(self, x):
output = self.embedding(x)
return output
m = Model()
output1 = m(x)
# compiled model
class Model(tf.keras.Model):
def __init__(self):
super(Model, self).__init__()
self.embedding = tf.keras.layers.Embedding(1, 1)
@tf.function(jit_compile=True)
def call(self, x):
output = self.embedding(x)
return output
m = Model()
output2 = m(x)
print(output1)
print(output2)Relevant log output
tf.Tensor([[0.]], shape=(1, 1), dtype=float32)
tf.Tensor([[-0.00567592]], shape=(1, 1), dtype=float32)