Fix fatal crash on invalid input rank in GPU linear algebra ops

[bodapatisaikrishna]

Description

Fixes a fatal crash (SIGABRT / Aborted (core dumped)) in GPU linear algebra operations (tf.linalg.det, tf.linalg.slogdet, tf.linalg.logdet, tf.linalg.cholesky, tf.linalg.inv, tf.linalg.solve, tf.linalg.qr) when given inputs with rank < 2 (such as a 0-D scalar). Reported in #76730.

Reproducer

import tensorflow as tf

with tf.device('/GPU:0'):
  # 0-D scalar input
  x = tf.zeros([])
  tf.linalg.det(x)

Before this change:
The process is killed immediately by SIGABRT:

tensorflow/core/framework/tensor_shape.cc:361: Check failed: d >= 0 (0 vs. -1)
Aborted (core dumped)

No Python exception is raised; the entire Python runtime / interpreter crashes.

After this change:
Cleanly raises tf.errors.InvalidArgumentError: Input must have rank >= 2, got 0 matching CPU behavior.

Root Cause

In GPU linear algebra kernels (DeterminantOpGpu, LogDeterminantOpGpu, CholeskyOpGpu, MatrixInverseOpGpu, MatrixSolveOpGpu, QrOpGpu), input dimension queries were executed before input rank validation:

// E.g. in DeterminantOpGpu / MatrixInverseOpGpu:
const Tensor& input = context->input(0);
const int ndims = input.dims();
const int64_t n = input.dim_size(ndims - 1);  // <-- ndims - 1 evaluates to -1 when ndims == 0

OP_REQUIRES_ASYNC(context, ndims >= 2, ...);

When input has rank 0 (scalar tf.zeros([])), ndims = 0, so ndims - 1 evaluates to -1. In tensor_shape.cc:361, TensorShapeBase::dim_size(int d) calls CHECK_GE(d, 0). When d = -1, the CHECK fails and calls abort(), killing the process.

Likewise:

• In QrOpGpu, input.dim_size(ndims - 2), input.dim_size(ndims - 1), and input.template flat_inner_dims<Scalar, 3>() were called before OP_REQUIRES_ASYNC(context, ndims >= 2, ...).
• In MatrixSolveOpGpu, input.dim_size(ndims - 1) and rhs.dim_size(ndims - 1) were called before OP_REQUIRES_ASYNC(context, ndims >= 2, ...) and OP_REQUIRES_ASYNC(context, rhs.dims() == ndims, ...).

In contrast, other GPU kernels (e.g. svd_op_gpu.cu.cc:356, lu_op_gpu.cu.cc:94, and self_adjoint_eig_v2_op_gpu.cc:51) and CPU kernels (via LinearAlgebraOp::AnalyzeInputs) validate ndims >= 2 before indexing into inner dimensions.

Solution

Hoist OP_REQUIRES_ASYNC(context, ndims >= 2, ...) (and OP_REQUIRES_ASYNC(context, rhs.dims() == ndims, ...) in MatrixSolveOpGpu) before accessing dim_size(ndims - 1) or dim_size(ndims - 2) across:

• tensorflow/core/kernels/linalg/determinant_op.cc (DeterminantOpGpu and LogDeterminantOpGpu)
• tensorflow/core/kernels/linalg/cholesky_op_gpu.cu.cc (CholeskyOpGpu)
• tensorflow/core/kernels/linalg/matrix_inverse_op.cc (MatrixInverseOpGpu)
• tensorflow/core/kernels/linalg/matrix_solve_op.cc (MatrixSolveOpGpu)
• tensorflow/core/kernels/linalg/qr_op_impl.h (QrOpGpu)

This ensures rank validation runs first, returning a standard InvalidArgumentError instead of terminating the process with SIGABRT.

Testing

Added testInvalidRank tests across graph and eager modes with use_gpu=True covering rank 0 (scalar) and rank 1 (vector) inputs across:

• tensorflow/python/kernel_tests/linalg/determinant_op_test.py
• tensorflow/python/kernel_tests/linalg/cholesky_op_test.py
• tensorflow/python/kernel_tests/linalg/matrix_inverse_op_test.py
• tensorflow/python/kernel_tests/linalg/matrix_solve_op_test.py
• tensorflow/python/kernel_tests/linalg/qr_op_test.py

All lines adhere to Google C++ and Python style guidelines (<= 80 characters per line).

Fixes #76730.

[gemini-code-assist]

Code Review

This pull request updates several GPU linear algebra operations (Cholesky, Determinant, Matrix Inverse, Matrix Solve, and QR) to perform input rank validation before accessing inner dimensions, preventing potential out-of-bounds errors. It also adds corresponding unit tests to verify behavior with invalid ranks. The feedback identifies a missing import of 'errors_impl' in 'qr_op_test.py' that will cause a 'NameError' when running the new test.

[gemini-code-assist]

The test testInvalidRank uses errors_impl.InvalidArgumentError, but errors_impl is not imported in this file. This will cause a NameError when running the test. Please import errors_impl from tensorflow.python.framework.

Suggested change

	from tensorflow.python.ops import gen_linalg_ops
	from tensorflow.python.framework import errors_impl
	from tensorflow.python.ops import gen_linalg_ops

[bodapatisaikrishna]

errors_impl is already imported at line 23 of this file (from tensorflow.python.framework import errors_impl) alongside the framework imports, and is used by existing tests such as testWrongDimensions on line 49.

[Abhirup0]

The change looks solid and clean. Verified that hoisting OP_REQUIRES_ASYNC(ndims >= 2) across CholeskyOpGpu, DeterminantOpGpu, LogDeterminantOpGpu, MatrixInverseOpGpu, MatrixSolveOpGpu, and QrOpGpu prevents the fatal CHECK_GE failure in TensorShapeBase::dim_size when inputs are scalars or 1-D vectors, matching the CPU behavior in LinearAlgebraOp::AnalyzeInputs.

One minor suggestion for matrix_solve_op_test.py:
In testInvalidRank, passing fn(val, val) tests cases where both inputs have rank < 2. To also exercise the second hoisted check (rhs.dims() == ndims), it would be good to include an asymmetric case where matrix has rank 2 but rhs has rank < 2:

valid_matrix = constant_op.constant(np.eye(2, dtype=np.float32))
for bad_shape in ([], [2]):
  bad_rhs = constant_op.constant(np.zeros(bad_shape, dtype=np.float32))
  with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
    with test_util.use_gpu():
      self.evaluate(fn(valid_matrix, bad_rhs))

[bodapatisaikrishna]

Thank you @Abhirup0 for the review and great catch on exercising the second hoisted check (rhs.dims() == ndims)!

Added the asymmetric test case where matrix has rank 2 and rhs has rank < 2 in commit 0d6dd8e:

        with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
          with test_util.use_gpu():
            self.evaluate(fn(valid_matrix, bad_val))

[dmiltr3]

Thank you for the fix. This perfectly identifies and resolves the out-of-bounds dim_size(ndims - 1) crash by properly hoisting OP_REQUIRES_ASYNC(context, ndims >= 2, ...) ahead of tensor dimension inquiries. The change is extremely targeted and avoids introducing any python-level overhead.

We require one testing update before approval.

1. (P1) Expand Data Type Coverage in Invalid Rank Tests

File: tensorflow/python/kernel_tests/linalg/*_op_test.py

The newly added testInvalidRank blocks correctly cover bad_shape boundary values [] and [2]. However, per our testing guidelines, we require unit tests to cover the Cartesian product of all registered dtypes for the ops. Could you please expand the test coverage across all test files to iterate over float and complex dtypes?

Example snippet:

      for bad_shape in ([], [2]):
        for dtype in (np.float32, np.float64, np.complex64, np.complex128):
          val = constant_op.constant(np.zeros(bad_shape, dtype=dtype))
          with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
            with test_util.use_gpu():
              self.evaluate(fn(val))

2. (P2: Nit) Missing space in existing error string

File: tensorflow/core/kernels/linalg/determinant_op.cc (Lines 139 and 278)

While moving the lines related to the square matrices check, there is a pre-existing missing space in the concat string "Input matrices must be square, got" which produces run-on strings. A space after "got " would be appreciated while you form the new testing commits.

[bodapatisaikrishna]

Added the dtype loop across all test files for float32, float64, complex64, and complex128, and fixed the spacing in the error string.

[djhdusjbd-eng]

I value this offer; could you share additional details?

…

On Mon, Sep 28, 2026 at 2:08 PM bodapatisaikrishna ***@***.***> wrote: *bodapatisaikrishna* left a comment (tensorflow/tensorflow#128068) <#128068 (comment)> Added the dtype loop across all test files for float32, float64, complex64, and complex128, and fixed the spacing in the error string. — Reply to this email directly, view it on GitHub <#128068?email_source=notifications&email_token=B4CX4CFZODIHXSX5TRRQ7MD5RJPGPA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKOBXGAZTENRXG422M4TFMFZW63VKON2WE43DOJUWEZLEUVSXMZLOOSWGM33PORSXEX3DNRUWG2Y#issuecomment-5870326775>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/B4CX4CGUDGYYFIGN3RSMEAT5RJPGPAVCNFSNUABEKJSXA33TNF2G64TZHM2DKNZRG4ZDKMB3JFZXG5LFHM2TKOBUGEYDCMJWGWQXMAQ> . Triage notifications, keep track of coding agent tasks and review pull requests on the go with GitHub Mobile for iOS <https://github.com/notifications/mobile/ios/B4CX4CAMQXT6S4PLWPL4EQ35RJPGPA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKOBXGAZTENRXG422M4TFMFZW63VKON2WE43DOJUWEZLEUVSXMZLOOSVGM33PORSXEX3JN5ZQ> and Android <https://github.com/notifications/mobile/android/B4CX4CH22HHQSGJQKJLHAGD5RJPGPA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKOBXGAZTENRXG422M4TFMFZW63VKON2WE43DOJUWEZLEUVSXMZLOOSXGM33PORSXEX3BNZSHE33JMQ>. Download it today! You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[dmiltr3]

The proposed fix cleanly resolves fatal crashes caused by out-of-bounds dimension accesses on tensors with rank < 2 in GPU linear algebra ops (CholeskyOpGpu, DeterminantOpGpu, LogDeterminantOpGpu, MatrixInverseOpGpu, MatrixSolveOpGpu, QrOpGpu). Hoisting the OP_REQUIRES_ASYNC(context, ndims >= 2, ...) validation check before indexing dim_size(ndims - 1) or invoking flat_inner_dims eliminates undefined behavior and process termination while introducing zero overhead on valid inputs.

The expanded test coverage across the Cartesian product of standard floating-point and complex dtypes (np.float32, np.float64, np.complex64, np.complex128) for both scalar and 1D vector inputs thoroughly verifies both eager and graph modes across public APIs and raw ops.

Optional Polish Suggestion (Non-blocking):

• Intra-module error string consistency:
  In tensorflow/core/kernels/linalg/determinant_op.cc:139, the non-square matrix error message was updated to include a trailing space:
  "Input matrices must be square, got ".
  In tensorflow/core/kernels/linalg/cholesky_op_gpu.cu.cc:114 and tensorflow/core/kernels/linalg/matrix_inverse_op.cc:153, the existing string is "Input matrices must be squares, got" (using plural "squares" and missing a space after "got").
  While this is pre-existing code outside the direct scope of the rank fix, you may optionally standardize them to "Input matrices must be square, got " for consistency across the linear algebra module.

[bodapatisaikrishna]

Standardized the error strings in cholesky_op_gpu.cu.cc and matrix_inverse_op.cc as well. Thanks for the review!

[dmiltr3]

This is a clean and surgical fix that correctly identifies and prevents fatal SIGABRT / Check failed: d >= 0 crashes in GPU linear algebra ops without introducing runtime overhead. Placing the bounds validation before inner dimension computations and memory allocations inside the C++ OpKernel is the canonical and safest approach.

The tests adequately cover the $N=0$ scalar and $N=1$ vector edge cases across the respective op variations.

I have just one minor nit for intra-kernel consistency:

• In tensorflow/core/kernels/linalg/matrix_solve_op.cc:150, consider tweaking "Input matrices must be squares, got " to "Input matrices must be square, got ". This would make it identical to the grammar cleanups you already performed in cholesky_op_gpu, determinant_op, and matrix_inverse_op.

Otherwise, this change looks excellent and is safe to merge.

[bodapatisaikrishna]

Updated the error string in matrix_solve_op.cc as well. Thank you for the review and approval!

[dmiltr3]

Summary

Thank you for this high-quality fix. This change addresses issue #76730 where passing rank < 2 inputs (such as 0-D scalar tensors) to GPU linear algebra operations triggered a process termination (SIGABRT due to a failed CHECK_GE(d, 0) in TensorShapeBase::dim_size).

By hoisting the rank validation checks before indexing into inner dimensions (input.dim_size(ndims - 1) / ndims - 2) and before computing inner flat dimensions, these operations now cleanly raise standard tf.errors.InvalidArgumentError exceptions consistent with CPU behavior.

Key Strengths

• Correct Invariant Sequencing: Moving OP_REQUIRES_ASYNC(context, ndims >= 2, ...) and rhs.dims() == ndims ahead of dimension queries eliminates the negative index calculation (d = -1 or -2) on rank-0 and rank-1 inputs.
• Kernel-Level Validation: Placing validation directly in ComputeAsync ensures full protection across eager execution, GraphDef execution, and direct gen_linalg_ops raw op dispatch.
• Minimal and Targeted Diff: The changes across tensorflow/core/kernels/linalg/cholesky_op_gpu.cu.cc, tensorflow/core/kernels/linalg/determinant_op.cc, tensorflow/core/kernels/linalg/matrix_inverse_op.cc, tensorflow/core/kernels/linalg/matrix_solve_op.cc, and tensorflow/core/kernels/linalg/qr_op_impl.h are concise, clean, and zero-cost on the steady-state fast path.
• Thorough Test Coverage: The added testInvalidRank suites in tensorflow/python/kernel_tests/linalg/ comprehensively test both rank-0 ([]) and rank-1 ([2]) inputs across graph and eager modes with use_gpu=True, verifying all registered dtypes (float32, float64, complex64, complex128) on both high-level APIs and raw op kernels.

The change is approved for integration.

[dmiltr3]

Thank you for your pull request and for addressing the fatal crash on rank < 2 inputs in GPU linear algebra ops.

During our internal code review and automated test suite verification, two test coverage enhancements were highlighted to ensure complete testing of the C++ GPU kernel implementations:

---

1. Converse Asymmetric Rank Pairing in matrix_solve_op_test.py

In tensorflow/python/kernel_tests/linalg/matrix_solve_op_test.py, testInvalidRank tests both symmetric invalid ranks fn(bad_val, bad_val) and asymmetric fn(valid_matrix, bad_val) (valid LHS, invalid RHS).

To be exhaustive, please also test the converse asymmetric pairing fn(bad_val, valid_matrix) (invalid LHS with rank < 2, valid RHS with rank >= 2).

Suggested Diff:

--- a/tensorflow/python/kernel_tests/linalg/matrix_solve_op_test.py
+++ b/tensorflow/python/kernel_tests/linalg/matrix_solve_op_test.py
@@ -133,3 +133,8 @@ class MatrixSolveOpTest(test.TestCase):
             with test_util.use_gpu():
               self.evaluate(fn(valid_matrix, bad_val))
+          with self.assertRaises(
+              (ValueError, errors_impl.InvalidArgumentError)
+          ):
+            with test_util.use_gpu():
+              self.evaluate(fn(bad_val, valid_matrix))

---

2. Exercise Dynamic Shapes in Graph Mode via array_ops.placeholder_with_default

In graph mode, tensors constructed with constant_op.constant(np.zeros(bad_shape)) have statically known shapes at graph construction time. TensorFlow's C++ shape inference intercepts invalid ranks during graph building and raises a ValueError before execution. As a result, in graph mode, execution never reaches the GPU kernel's ComputeAsync. While eager mode exercises the C++ kernel directly, graph-mode runtime kernel execution should also be validated.

To verify the GPU kernel at runtime in graph mode as well, please test dynamic shapes with un-inferred rank using array_ops.placeholder_with_default(..., shape=None).

Example Pattern:

          # Test with static rank
          with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
            with test_util.use_gpu():
              self.evaluate(fn(val))
          # Test with dynamic rank to exercise the C++ GPU kernel at runtime in graph mode
          val_dyn = array_ops.placeholder_with_default(val, shape=None)
          with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
            with test_util.use_gpu():
              self.evaluate(fn(val_dyn))

Please apply this dynamic shape check to testInvalidRank in the relevant linear algebra tests:

• tensorflow/python/kernel_tests/linalg/cholesky_op_test.py
• tensorflow/python/kernel_tests/linalg/determinant_op_test.py
• tensorflow/python/kernel_tests/linalg/matrix_inverse_op_test.py
• tensorflow/python/kernel_tests/linalg/matrix_solve_op_test.py
• tensorflow/python/kernel_tests/linalg/qr_op_test.py

Once these test updates are pushed, we will proceed with running the automated test suite and completing integration. Thank you!

[bodapatisaikrishna]

Added dynamic shape tests via placeholder_with_default across the linear algebra suites and added the converse asymmetric test for matrix_solve.

[dmiltr3]

Thank you for this fix! The implementation cleanly resolves the process termination issue caused by unchecked dimension validations across the GPU linear algebra kernels by correctly hoisting the checks ahead of evaluation.

We completely agree with your solution and all reviewers have signed off on the C++ kernel logic.

However, during our automated evaluation pipelines, two of the test environments threw execution failures:
tensorflow/python/kernel_tests/linalg/determinant_op_test.py and tensorflow/python/kernel_tests/linalg/matrix_inverse_op_test.py are both failing with NameError: name 'errors_impl' is not defined.

    with self.assertRaisesRegex(
        (ValueError, errors_impl.InvalidArgumentError),
        "Input must have rank >= 2, got 0"):

It appears errors_impl is missing from the imports at the top of these specific test files. Please add from tensorflow.python.framework import errors_impl to the test files where it was omitted.

Once the imports are fixed, this is fully approved and we will merge it!

[bodapatisaikrishna]

Updated the imports and added the missing errors dependencies in the BUILD file so the test environments pick them up properly.

[dmiltr3]

Thank you for addressing the test suite imports! Hoisting the rank validation ahead of dimension indexing across all GPU linear algebra kernels cleanly resolves the fatal crash, and the expanded test coverage is thorough.

Before this pull request can be merged, please resolve one remaining issue causing the PyLint CI check to fail:

Remove unused errors import and build dependencies (Required)

In tensorflow/python/kernel_tests/linalg/determinant_op_test.py (line 21) and tensorflow/python/kernel_tests/linalg/matrix_inverse_op_test.py (line 22), errors was imported alongside errors_impl:

from tensorflow.python.framework import errors
from tensorflow.python.framework import errors_impl

Because only errors_impl.InvalidArgumentError is used in the tests, PyLint fails with W0611: Unused errors imported from tensorflow.python.framework (unused-import):

tensorflow/python/kernel_tests/linalg/determinant_op_test.py:21:0: W0611: Unused errors imported from tensorflow.python.framework (unused-import)
tensorflow/python/kernel_tests/linalg/matrix_inverse_op_test.py:22:0: W0611: Unused errors imported from tensorflow.python.framework (unused-import)

Please remove the unused errors import from both test files:

# In determinant_op_test.py and matrix_inverse_op_test.py:
-from tensorflow.python.framework import errors
 from tensorflow.python.framework import errors_impl

And in tensorflow/python/kernel_tests/linalg/BUILD, remove the unused "//tensorflow/python/framework:errors" dependency from determinant_op_test and matrix_inverse_op_test:

# In tensorflow/python/kernel_tests/linalg/BUILD:
 cuda_py_strict_test(
     name = "determinant_op_test",
     deps = [
         "//tensorflow/python/client:session",
         "//tensorflow/python/framework:constant_op",
-        "//tensorflow/python/framework:errors",
         "//tensorflow/python/framework:for_generated_wrappers",
...
 cuda_py_strict_test(
     name = "matrix_inverse_op_test",
     deps = [
         "//tensorflow/python/client:session",
         "//tensorflow/python/framework:constant_op",
-        "//tensorflow/python/framework:errors",
         "//tensorflow/python/framework:for_generated_wrappers",

This will match sibling tests (cholesky_op_test, matrix_solve_op_test, qr_op_test) which only depend on and import errors_impl, and will get PyLint passing cleanly. Once this is removed, this PR will be approved!

[bodapatisaikrishna]

Removed the unused errors import and build target dependencies so PyLint passes cleanly.

[dmiltr3]

Thank you for addressing the remaining PyLint issues perfectly. Your updates tracking the input dimensionality correctly above the evaluation loops gracefully prevents those SIGABRT crashes across the entire suite of GPU linear algebra implementations while adhering strictly to performance constraints.

All CI checks are passing beautifully. Thank you incredibly for seeing this through. The PR is fully approved.