[XLA] crashes when compiling tf.raw_ops.AssignVariableOp #88038
Description
Issue type
Bug
Have you reproduced the bug with TensorFlow Nightly?
Yes
Source
source
TensorFlow version
nightly 20250223
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?
tf.raw_ops.AssignVariableOp should pass the compliation
Standalone code to reproduce the issue
import os
import tensorflow
import tensorflow as tf
import numpy as np
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
class VariableModel(tf.keras.Model):
def __init__(self):
super(VariableModel, self).__init__()
self.variable = tf.Variable(initial_value=0.0)
def call(self, x):
return tf.raw_ops.AssignVariableOp(resource=self.variable.handle, value=tf.add(self.variable, x))
model = VariableModel()
input_shape = [1]
x = tf.constant([1.0], shape=input_shape)
inputs = [x]
model(*inputs)
print("succeed on eager")
class VariableModel(tf.keras.Model):
def __init__(self):
super(VariableModel, self).__init__()
self.variable = tf.Variable(initial_value=0.0)
@tf.function(jit_compile=True)
def call(self, x):
return tf.raw_ops.AssignVariableOp(resource=self.variable.handle, value=tf.add(self.variable, x))
model = VariableModel()
model(*inputs)
print("succeed on XLA")Relevant log output
succeed on eager
ValueError: Shapes must be equal rank, but are 0 and 1 for '{{node AssignVariableOp}} = AssignVariableOp[dtype=DT_FLOAT, validate_shape=false](Add/ReadVariableOp/resource, Add)' with input shapes: [], [1].