tf.raw_ops.MatrixDiag crashes with Aborted (core dumped) with a 6D tensor as input on GPU #98349
Description
Issue type
Bug
Have you reproduced the bug with TensorFlow Nightly?
Yes
Source
binary
TensorFlow version
2.21.0-dev20250801
Custom code
Yes
OS platform and distribution
Linux Ubuntu 24.04
Mobile device
No response
Python version
3.12
Bazel version
No response
GCC/compiler version
No response
CUDA/cuDNN version
No response
GPU model and memory
No response
Current behavior?
Running tf.raw_ops.MatrixDiag with a 6D tensor as the input caused it to raise SIGABRT and crash with Aborted (core dumped) on GPU (it does not crash on CPU).
The crash can be reproduced with Tensorflow 2.19.0. Please take a look at the gist
Standalone code to reproduce the issue
import tensorflow as tf
import numpy as np
print("TensorFlow version:", tf.__version__) # TensorFlow version: 2.21.0-dev20250801
rng = np.random.default_rng(253)
diagonal = tf.constant(rng.uniform(0, 1, size=(9, 3, 70, 59, 22, 30)), dtype=tf.float32)
with tf.device("/GPU:0"):
output = tf.raw_ops.MatrixDiag(diagonal=diagonal)Relevant log output
F0000 00:00:1754429004.779581 2763615 gpu_launch_config.h:129] Check failed: work_element_count >= 0 (-2087069296 vs. 0)
*** Check failure stack trace: ***
@ 0x78c7c8a57874 absl::lts_20250127::log_internal::LogMessage::SendToLog()
@ 0x78c7c8a57206 absl::lts_20250127::log_internal::LogMessage::Flush()
@ 0x78c7c8a57e19 absl::lts_20250127::log_internal::LogMessageFatal::~LogMessageFatal()
@ 0x78c7a90cddf0 tensorflow::functor::MatrixDiag<>::Compute()
@ 0x78c7a90c300a tensorflow::MatrixDiagOp<>::Compute()
@ 0x78c7c8adf984 tensorflow::BaseGPUDevice::Compute()
@ 0x78c7c6ab5368 tensorflow::(anonymous namespace)::SingleThreadedExecutorImpl::Run()
@ 0x78c7c6a80ead tensorflow::FunctionLibraryRuntimeImpl::RunSync()
@ 0x78c7c6a8cae4 tensorflow::ProcessFunctionLibraryRuntime::RunMultiDeviceSync()
@ 0x78c7c6a92ccd tensorflow::ProcessFunctionLibraryRuntime::RunSync()
@ 0x78c7b9ac8ffd tensorflow::KernelAndDeviceFunc::Run()
@ 0x78c7b9a7610d tensorflow::EagerKernelExecute()
@ 0x78c7b9a7fd5b tensorflow::ExecuteNode::Run()
@ 0x78c7b9fde5bf tensorflow::EagerExecutor::SyncExecute()
@ 0x78c7b9a75a3c tensorflow::(anonymous namespace)::EagerLocalExecute()
@ 0x78c7b9a73268 tensorflow::DoEagerExecute()
@ 0x78c7b9a76e4f tensorflow::EagerExecute()
@ 0x78c7b96804e7 tensorflow::EagerOperation::Execute()
@ 0x78c7b9ac5987 tensorflow::CustomDeviceOpHandler::Execute()
@ 0x78c7b7044ef5 TFE_Execute
@ 0x78c789c98993 TFE_Py_FastPathExecute_C()
@ 0x78c77a2d81e3 pybind11::detail::argument_loader<>::call<>()
@ 0x78c77a2d811f pybind11::cpp_function::initialize<>()::{lambda()#1}::__invoke()
@ 0x78c77a2ad969 pybind11::cpp_function::dispatcher()
@ 0x581d4f (unknown)
@ 0x548f85 _PyObject_MakeTpCall
@ 0x5d6b2f _PyEval_EvalFrameDefault
@ 0x5d500b PyEval_EvalCode
@ 0x6081e2 (unknown)
@ 0x6b5033 (unknown)
@ 0x6b4d9a _PyRun_SimpleFileObject
@ 0x6b4bcf _PyRun_AnyFileObject
@ 0x6bcc35 Py_RunMain
@ 0x6bc71d Py_BytesMain
@ 0x78c7c9e2a1ca (unknown)
@ 0x78c7c9e2a28b __libc_start_main
@ 0x6575a5 _start
Aborted (core dumped)