Skip to content

Validate input rank in the tf2xla QR and SVD kernels - #126542

Merged
copybara-service[bot] merged 1 commit into
tensorflow:masterfrom
vishwakt:fix-tf2xla-qr-svd-rank
Sep 3, 2026
Merged

Validate input rank in the tf2xla QR and SVD kernels#126542
copybara-service[bot] merged 1 commit into
tensorflow:masterfrom
vishwakt:fix-tf2xla-qr-svd-rank

Conversation

@vishwakt

@vishwakt vishwakt commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

Fixes #110798.

Problem

The tf2xla kernels for Qr, Svd, and XlaSvd pass their input straight into the XLA builder libraries without validating its rank. The graph-level shape functions reject known-rank inputs below rank 2, but an input whose static rank is unknown, such as the reporter's StackPopV2 output or any unknown-rank placeholder, reaches the kernels unchecked at XLA compile time:

  • QROp::Compile calls xla::QrExplicit, which reads dimension rank - 2 = -1 and dies on the fatal ShapeUtil::GetDimensionNumber check (shape_util.cc:952, the check in the report).
  • SvdOp::Compile dies one step earlier, on its own input_shape.dim_size(dims() - 2) call, hitting the d >= 0 check in TensorShape::dim_size. This sibling was found while verifying the report and is fixed here as well; XlaSvdOp shares the same gap.

Both reproduce on XLA CPU with jit_compile=True (verified locally, exit 134 with the reported check-failure text for QR and the TensorShape variant for SVD). The neighboring Cholesky and SelfAdjointEig lowerings already reject rank-1 inputs with clean InvalidArgumentErrors, which confirms the intended behavior for this input class.

Fix

Validate rank >= 2 in all three kernels before any trailing dimensions are read, using the validation and message style of the neighboring MatrixSolve kernel ("Input must have rank >= 2, got shape ..."). Compilation now fails with a regular InvalidArgumentError instead of aborting the process.

Tests

testVectorInputRaisesError in compiler/tests/qr_op_test.py and compiler/tests/svd_op_test.py feeds a rank-1 value through an unknown-rank placeholder, which is exactly the path that bypasses graph-level shape inference; on unfixed builds this aborts the test process, and with the fix it asserts the InvalidArgumentError. BUILD dependencies updated for the strict-deps targets (kernel targets gain the absl status and strings deps used by the new checks, test targets gain //tensorflow/python/framework:errors).

The tf2xla kernels for Qr, Svd, and XlaSvd passed their input straight
to the XLA builder libraries without checking its rank. The graph-level
shape functions reject inputs of known rank below 2, but an input of
unknown static rank, for example one produced by StackPopV2, reaches
the kernels unchecked at compile time. QrExplicit then reads dimension
rank - 2 = -1 and dies on a fatal ShapeUtil check, and SvdOp dies on
the dim_size range check in TensorShape, aborting the process instead
of failing compilation.

Validate rank >= 2 in all three kernels before any trailing dimensions
are read, mirroring the validation and message style of the
neighboring MatrixSolve kernel and the errors already produced by the
Cholesky and SelfAdjointEig builder libraries.

The regression tests feed a rank-1 value through an unknown-rank
placeholder, which reproduces the abort on unfixed builds, and assert
that a regular InvalidArgumentError is raised instead.

Fixes tensorflow#110798
@google-ml-butler google-ml-butler Bot added the size:M CL Change Size: Medium label Sep 1, 2026
@nithyak0204
nithyak0204 requested a review from a team September 1, 2026 07:37
@google-ml-butler google-ml-butler Bot added the awaiting review Pull request awaiting review label Sep 1, 2026
@nithyak0204 nithyak0204 added the prtype:bugfix PR to fix a bug label Sep 1, 2026
@github-project-automation github-project-automation Bot moved this to Assigned Reviewer in PR Queue Sep 1, 2026
@dmiltr3
dmiltr3 self-requested a review September 1, 2026 13:07
@github-project-automation github-project-automation Bot moved this from Assigned Reviewer to Approved by Reviewer in PR Queue Sep 1, 2026
@google-ml-butler google-ml-butler Bot added kokoro:force-run Tests on submitted change ready to pull PR ready for merge process labels Sep 1, 2026
@nithyak0204 nithyak0204 removed awaiting review Pull request awaiting review kokoro:force-run Tests on submitted change labels Sep 2, 2026
This was referenced Sep 3, 2026
copybara-service Bot pushed a commit that referenced this pull request Sep 3, 2026
FUTURE_COPYBARA_INTEGRATE_REVIEW=#126542 from vishwakt:fix-tf2xla-qr-svd-rank ff2452d
PiperOrigin-RevId: 974183619
copybara-service Bot pushed a commit that referenced this pull request Sep 3, 2026
FUTURE_COPYBARA_INTEGRATE_REVIEW=#126542 from vishwakt:fix-tf2xla-qr-svd-rank ff2452d
PiperOrigin-RevId: 975474357
copybara-service Bot pushed a commit that referenced this pull request Sep 3, 2026
FUTURE_COPYBARA_INTEGRATE_REVIEW=#126542 from vishwakt:fix-tf2xla-qr-svd-rank ff2452d
PiperOrigin-RevId: 975474341
copybara-service Bot pushed a commit that referenced this pull request Sep 3, 2026
FUTURE_COPYBARA_INTEGRATE_REVIEW=#126542 from vishwakt:fix-tf2xla-qr-svd-rank ff2452d
PiperOrigin-RevId: 973270773
This was referenced Sep 3, 2026
@copybara-service
copybara-service Bot merged commit 91fd803 into tensorflow:master Sep 3, 2026
19 of 20 checks passed
@github-project-automation github-project-automation Bot moved this from Approved by Reviewer to Merged in PR Queue Sep 3, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

prtype:bugfix PR to fix a bug ready to pull PR ready for merge process size:M CL Change Size: Medium

Projects

Status: Merged

Development

Successfully merging this pull request may close these issues.

XLA CPU/GPU JIT CHECK-fails in QR (shape_util.cc:952 dimension_number >= 0) with StackV2 → StackPopV2 input

5 participants