XLA compilation (jit_compile=True) aborts with a fatal CHECK when compiling tf.where on the result of diagflat(tf.unique(...))

XLA compilation (jit_compile=True) aborts with a fatal internal assertion when compiling a valid TensorFlow graph that applies tf.where to the result of tf.experimental.numpy.diagflat(tf.unique(...)). The same computation executes successfully in eager mode. Instead of successfully compiling the graph or returning a structured TensorFlow/XLA error, the XLA compiler aborts the process during compilation with: ```Check failed: subshape->dimensions(dim_index) >= size (64 vs. 255)``` Here is the [gist](https://colab.research.google.com/gist/jasminetrail/03f92c692cd5bf59c9e9e97eeda1dc0b/issue_44941.ipynb) --- **Minimal code to reproduce** ```Python import traceback import tensorflow as tf tf.random.set_seed(12682) t20180 = tf.random.normal([32], dtype=tf.float32) def model_12682(t20180): t20191 = tf.unique(t20180, out_idx=tf.int32) t20192 = tf.experimental.numpy.sqrt(t20191) t20193 = tf.experimental.numpy.diagflat(t20192, k=0) t20195 = tf.where(t20193) return t20195 _args = (t20180,) print("=== EAGER ===") try: eager_out = model_12682(*_args) print("success") except Exception: print(traceback.format_exc()) print("\n=== XLA ===") try: xla_out = tf.function(model_12682, jit_compile=True)(*_args) print("success") except Exception: print(traceback.format_exc()) ``` ## Error logs ```Python WARNING: All log messages before absl::InitializeLog() is called are written to STDERR I0000 00:00:1782591553.593787 1055320 port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. I0000 00:00:1782591553.594171 1055320 cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used. I0000 00:00:1782591553.640333 1055320 cpu_feature_guard.cc:227] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations. To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags. WARNING: All log messages before absl::InitializeLog() is called are written to STDERR I0000 00:00:1782591554.808092 1055320 port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`. I0000 00:00:1782591554.808796 1055320 cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used. E0000 00:00:1782591554.920142 1055320 cuda_platform.cc:52] failed call to cuInit: INTERNAL: CUDA error: Failed call to cuInit: UNKNOWN ERROR (303) === EAGER === success === XLA === I0000 00:00:1782591555.124913 1055320 service.cc:153] XLA service 0x64cb482a6630 initialized for platform Host (this does not guarantee that XLA will be used). Devices: I0000 00:00:1782591555.124955 1055320 service.cc:161] StreamExecutor [0]: Host, Default Version (Driver: 0.0.0; Runtime: 0.0.0; Toolkit: 0.0.0; DNN: 0.0.0) W0000 00:00:1782591555.127181 1055320 assert_op.cc:39] Ignoring Assert operator Assert/Assert I0000 00:00:1782591555.131219 1055320 dump_mlir_util.cc:269] disabling MLIR crash reproducer, set env var `MLIR_CRASH_REPRODUCER_DIRECTORY` to enable. F0000 00:00:1782591555.137310 1055320 literal.cc:758] Check failed: subshape->dimensions(dim_index) >= size (64 vs. 255) *** Check failure stack trace: *** @ 0x76e3ec71d314 absl::lts_20250814::log_internal::LogMessage::SendToLog() @ 0x76e3ec71d296 absl::lts_20250814::log_internal::LogMessage::Flush() @ 0x76e3eb451d9d xla::MutableLiteralBase::SetDynamicSize() @ 0x76e3eb455296 xla::LiteralBase::Broadcast() @ 0x76e3ea9bbec9 xla::HloEvaluator::HandleBroadcast() @ 0x76e3eb36e540 xla::HloInstruction::Visit<>() @ 0x76e3ea9ae46b xla::HloEvaluator::EvaluateInternal() @ 0x76e3ea9ad9d8 xla::HloEvaluator::Evaluate() @ 0x76e3dc7a8496 xla::(anonymous namespace)::HloProtoEvaluator::Evaluate() @ 0x76e3dc7bbd2b std::_Function_handler<>::_M_invoke() @ 0x76e3dc7a09e9 xla::(anonymous namespace)::PostorderDFSVisitor::PostOrderDFSVisit() @ 0x76e3dc7a94b1 xla::ValueInference::AnalyzeConstant() @ 0x76e3e2cd2bcb xla::SetDimensionSizeWithRebound() @ 0x76e3e2cd1c89 tensorflow::(anonymous namespace)::WhereOp::Compile() @ 0x76e3e2d87ecc tensorflow::XlaOpKernel::Compute() @ 0x76e3dc790d87 tensorflow::XlaCompilationDevice::Compute() @ 0x76e3e2d6aa03 tensorflow::GraphCompiler::Compile() @ 0x76e3e2d77a34 tensorflow::(anonymous namespace)::ExecuteGraph() @ 0x76e3e2d710fa tensorflow::XlaCompiler::CompileGraph() @ 0x76e3e2d72dc3 tensorflow::XlaCompiler::CompileFunction() @ 0x76e3e2cf7d4f tensorflow::TfGraphToHloCompiler::Compile() @ 0x76e3e28dc01d tensorflow::DeviceCompiler<>::CompileStrict() @ 0x76e3e28dabc0 tensorflow::DeviceCompiler<>::CompileImpl() @ 0x76e3e28da1b5 tensorflow::DeviceCompiler<>::CompileIfNeeded() @ 0x76e3e28d3bbb tensorflow::(anonymous namespace)::CompileToLocalExecutable() @ 0x76e3e28d24ea tensorflow::XlaLocalLaunchBase::ComputeAsync() @ 0x76e3ec128c8d tensorflow::AsyncOpKernel::Compute() @ 0x76e3ea260e80 tensorflow::ThreadPoolDevice::Compute() @ 0x76e3e2df03d0 tensorflow::KernelAndDeviceOp::Run() @ 0x76e3e2da00fb tensorflow::EagerKernelExecute() @ 0x76e3e2da981e tensorflow::ExecuteNode::Run() @ 0x76e3e336f7b4 tensorflow::EagerExecutor::SyncExecute() @ 0x76e3e2d9faa1 tensorflow::(anonymous namespace)::EagerLocalExecute() @ 0x76e3e2d9d1d0 tensorflow::DoEagerExecute() @ 0x76e3e2da0b2b tensorflow::EagerExecute() @ 0x76e3e28ce247 tensorflow::EagerOperation::Execute() @ 0x76e3e2dedbb7 tensorflow::CustomDeviceOpHandler::Execute() @ 0x76e3dff004e5 TFE_Execute @ 0x76e3afc8e359 TFE_Py_ExecuteCancelable() @ 0x76e3aeb9d086 tensorflow::TFE_Py_ExecuteCancelable_wrapper() @ 0x76e3aebdc54d pybind11::cpp_function::initialize<>()::{lambda()#1}::__invoke() @ 0x76e3aebb3fdd pybind11::cpp_function::dispatcher() @ 0x64cb3bd7ae32 (unknown) Aborted (core dumped) ``` Versions: - Tensorflow: 2.21.0 - Python: 3.10