XLA compilation fails with Unimplemented padding for instruction: bitcast-convert when applying tf.bitcast to the output of tf.raw_ops.Where

[jasminetrail]

tf.function(jit_compile=True) fails to compile a graph that executes successfully in eager mode when tf.bitcast is applied to the output of tf.raw_ops.Where. The reproducer uses Where on a complex128 tensor to produce an integer index tensor, followed by tf.bitcast from int64 to float32. Eager execution completes successfully, but XLA compilation fails with:

UnimplementedError: Unimplemented padding for instruction: bitcast-convert(%Where...)

The failure occurs during XLA compilation and appears to be triggered by tf.bitcast on the dynamically sized output of tf.raw_ops.Where. Since the same graph executes correctly in eager mode, this indicates an inconsistency in XLA's lowering of bitcast-convert for Where outputs. Here is the gist

Minimal Code

import traceback
import tensorflow as tf

tf.random.set_seed(9776)
t95045 = tf.complex(tf.random.normal([16], dtype=tf.float64), tf.random.normal([16], dtype=tf.float64))

def model_9776(t95045):
    t95047 = tf.raw_ops.Where(condition=t95045)
    t95048 = tf.bitcast(t95047, tf.float32)
    return t95048

_args = (t95045,)

print("=== EAGER ===")
try:
    eager_out = model_9776(*_args)
    print("Eager Output Succeeded")
except Exception:
    print(traceback.format_exc())

print("\n=== XLA ===")
try:
    xla_out = tf.function(model_9776, jit_compile=True)(*_args)
    print("Compiler Output Succeeded")
except Exception:
    print(traceback.format_exc())

Error logs

=== EAGER ===
Eager Output Succeeded

=== XLA ===
Traceback (most recent call last):
  File "/tmp/ipykernel_1541/2987141401.py", line 23, in <cell line: 0>
    xla_out = tf.function(model_9776, jit_compile=True)(*_args)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tensorflow/python/util/traceback_utils.py", line 167, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/usr/local/lib/python3.12/dist-packages/tensorflow/python/eager/execute.py", line 59, in quick_execute
    except TypeError as e:
tensorflow.python.framework.errors_impl.UnimplementedError: Unimplemented padding for instruction: %Bitcast.1 = f32[16,1,2]{2,1,0} bitcast-convert(%Where.143), metadata={op_type="Bitcast" op_name="Bitcast" source_file="/usr/local/lib/python3.12/dist-packages/tensorflow/python/framework/ops.py" source_line=1221} [Op:__inference_model_9776_61]

Version

Tensorflow: 2.21.0