[Ops][addmv] Validate self shape and dimension in torch.addmv by Kaap10 · Pull Request #6011 · google-pytorch/torch_tpu

Kaap10 · 2026-10-29T10:57:53Z

Description

torch.addmv(self, mat, vec) validates that mat.dim() == 2, vec.dim() == 1, and mat.size(1) == vec.size(0) in csrc/ops/addmv/addmv_aten_kernels.cc. However, ValidateAddmvInputs performed no dimension or shape validation on self:

It omitted validating that self is a scalar or vector (self.dim() <= 1).
When self is 1-D, it omitted validating that its length matches mat.size(0) (or is 1 for singleton broadcast).

This caused two issues:

Silent execution with beta=0: Addmv executes the if (beta.IsZero()) branch which completely drops self and delegates to BuildAddmvShloBetaZero. Because self was never validated in ValidateAddmvInputs, invalid inputs (e.g. 2D/3D tensors or mismatched 1D lengths) silently succeeded on TPU, whereas PyTorch CPU/CUDA raises RuntimeError.
Unhandled StableHLO failure with beta != 0: Incompatible shapes failed deep inside StableHLO lowering with failed to broadcast tensor: ... instead of a standard PyTorch RuntimeError.

This PR:

Adds self.dim() <= 1 validation in ValidateAddmvInputs.
Adds self.size(0) == 1 || self.size(0) == mat.size(0) length validation for 1-D self in ValidateAddmvInputs.
Adds comprehensive error tests in tests/errors_test.py (test_addmv_not_a_vector_or_scalar and test_addmv_size_mismatch) covering both beta=1 and beta=0.

Test Plan

Verified Python formatting and linting: ruff check tests/errors_test.py (all checks passed with 0 errors).
Verified Google C++ formatting: all modified lines in csrc/ops/addmv/addmv_aten_kernels.cc strictly adhere to <= 80 columns.
Verified test coverage for both rank and length checks under default beta=1 and beta=0.

github-actions · 2026-10-29T10:58:02Z

Thanks for contributing to TorchTPU, @Kaap10! 🔥

To ensure smooth code reviews, we kindly ask authors and reviewers to follow the process documented in https://github.com/google-pytorch/torch_tpu/blob/main/CONTRIBUTING.md#code-reviews.

Happy hacking!

ghpvnist

Thanks for the contribution! Added a few comments.

Refine the size mismatch error message in addmv_aten_kernels.cc to include all operand shapes, and update reviewer attribution to gunhyun in tests/errors_test.py per review comments on google-pytorch#6011.

Kaap10 · 2026-10-30T02:45:46Z

Thanks for the review - @ghpvnist

I have addressed all the comments:

Updated the TPU error message in addmv_aten_kernels.cc to display all operand shapes: expected arguments shapes [5], [2, 3], [3] to have compatible shapes, got 5 vs 2 which is not compatible.
Updated message_reviewed_by="gunhyun" across all newly added assertions in tests/errors_test.py.
Verified with ruff check and confirmed all checks are passing.

Ready for another look whenever you have time!

ghpvnist

Thanks!

zhanyong-wan

Thanks!

zhanyong-wan · 2026-11-03T01:07:05Z

+
+    with et.assert_raises_message(
+        RuntimeError,
+        tpu="""addmv(): expected the first argument to be a vector or scalar (0D or 1D tensor), got 2D tensor""",


Can we swap "vector" and "scalar" to match the order of "0D or 1D"?

Done, swapped to "scalar or vector".

zhanyong-wan · 2026-11-03T01:08:22Z

+
+    with et.assert_raises_message(
+        RuntimeError,
+        tpu="""addmv(): expected arguments shapes [5], [2, 3], [3] to have compatible shapes, got 5 vs 2 which is not compatible""",


Grammar nit: "arguments shapes" => "argument shapes"

Done, updated to "argument shapes".

zhanyong-wan · 2026-11-03T01:09:04Z

+
+    with et.assert_raises_message(
+        RuntimeError,
+        tpu="""addmv(): expected arguments shapes [5], [2, 3], [3] to have compatible shapes, got 5 vs 2 which is not compatible""",


Nit: simplify "have compatible shapes" to "be compatible"?

Then we can drop " which is not compatible" as it's obvious.

Done, simplified to "to be compatible, got 5 vs 2".

zhanyong-wan

Please fix the CI errors.

Kaap10 · 2026-11-03T06:41:18Z

Please fix the CI errors.

Fixed.

Before:

        << " to be compatible, got " << self.size(0) << " vs "
        << mat.size(0);

After:

        << " to be compatible, got " << self.size(0) << " vs " << mat.size(0);

Verified locally with clang-format.

zhanyong-wan · 2026-11-10T18:38:50Z

@ghpvnist , the imported CL fails in internal testing. Could you fix up cl/1010873988 to get this landed? Thanks.

ghpvnist · 2026-11-10T21:52:44Z

+        RuntimeError,
+        tpu="""addmv(): expected the first argument to be a scalar or vector (0D or 1D tensor), got 2D tensor""",
+        gpu="""vector + matrix @ vector expected, got 2, 2, 1""",
+        message_reviewed_by="gunhyun",


There's a rebase conflict here and recently there was an api change to remove the message_reviewd_by parameter. Could you rebase to head and remove this field?

Done, rebased onto latest head and removed the message_reviewed_by parameter from the error tests.
Also updated the error macro registry in tests/error_coverage_data.yaml.

Fixes google-pytorch#6010 Validate that self is a scalar or vector (dim <= 1) and that when dim == 1, its length matches mat.size(0) or is 1 in ValidateAddmvInputs before dispatching to Addmv. This prevents silent execution on invalid shapes when beta=0 and unhandled StableHLO broadcast failures when beta!=0. Add test coverage in tests/errors_test.py.

Refine the size mismatch error message in addmv_aten_kernels.cc to include all operand shapes, and update reviewer attribution to gunhyun in tests/errors_test.py per review comments on google-pytorch#6011.

github-actions Bot added the stat:awaiting-reviewers label Oct 29, 2026

Kaap10 mentioned this pull request Oct 29, 2026

[Ops][addmv] Missing shape and dimension validation for self in torch.addmv #6010

Closed

benawilson requested a review from ghpvnist October 29, 2026 11:30

ghpvnist requested changes Oct 30, 2026

View reviewed changes

Comment thread tests/errors_test.py Resolved

Comment thread tests/errors_test.py Resolved

Comment thread tests/errors_test.py Resolved

Comment thread tests/errors_test.py Resolved

Comment thread tests/errors_test.py Resolved

Kaap10 requested a review from ghpvnist October 30, 2026 02:46

ghpvnist approved these changes Oct 30, 2026

View reviewed changes

jonathan1920 added stat:awaiting-requester and removed stat:awaiting-reviewers labels Oct 30, 2026

ghpvnist added pull ready stat:awaiting-requester and removed stat:awaiting-requester pull ready labels Oct 30, 2026

ghpvnist assigned zhanyong-wan Oct 31, 2026

github-actions Bot added stat:awaiting-reviewers and removed stat:awaiting-requester labels Oct 31, 2026

ghpvnist requested a review from zhanyong-wan October 31, 2026 00:14

zhanyong-wan requested changes Nov 3, 2026

View reviewed changes

Kaap10 requested a review from zhanyong-wan November 3, 2026 03:34

zhanyong-wan approved these changes Nov 3, 2026

View reviewed changes

zhanyong-wan requested changes Nov 3, 2026

View reviewed changes

github-actions Bot added stat:awaiting-requester and removed stat:awaiting-reviewers labels Nov 3, 2026

Kaap10 force-pushed the fix-addmv-self-validation branch from c98da2a to 0092530 Compare November 3, 2026 06:36

github-actions Bot added stat:awaiting-reviewers and removed stat:awaiting-requester labels Nov 3, 2026

zhanyong-wan approved these changes Nov 3, 2026

View reviewed changes

zhanyong-wan added the pull ready label Nov 3, 2026

github-actions Bot added stat:awaiting-requester and removed stat:awaiting-reviewers labels Nov 9, 2026

zhanyong-wan requested a review from ghpvnist November 10, 2026 18:38

github-actions Bot added stat:awaiting-reviewers and removed stat:awaiting-requester labels Nov 10, 2026

ghpvnist requested changes Nov 10, 2026

View reviewed changes

github-actions Bot added stat:awaiting-requester and removed stat:awaiting-reviewers labels Nov 10, 2026

Kaap10 added 4 commits November 11, 2026 07:51

[Ops][addmv] Refine error message and update reviewer attribution

36e15a9

Refine the size mismatch error message in addmv_aten_kernels.cc to include all operand shapes, and update reviewer attribution to gunhyun in tests/errors_test.py per review comments on google-pytorch#6011.

[Ops][addmv] Address review feedback on error message wording

2b4e504

[Ops][addmv] Remove message_reviewed_by and update error coverage

64b04d8

Kaap10 force-pushed the fix-addmv-self-validation branch from 0092530 to 64b04d8 Compare November 11, 2026 02:38

Kaap10 requested a review from ghpvnist November 11, 2026 02:40

github-actions Bot added stat:awaiting-reviewers and removed stat:awaiting-requester labels Nov 11, 2026

copybara-service Bot closed this in 9218ffd Nov 11, 2026

Conversation

Kaap10 commented Oct 29, 2026

Description

Test Plan

Uh oh!

github-actions Bot commented Oct 29, 2026

Uh oh!

ghpvnist left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Kaap10 commented Oct 30, 2026

Uh oh!

ghpvnist left a comment

Choose a reason for hiding this comment

Uh oh!

zhanyong-wan left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zhanyong-wan left a comment

Choose a reason for hiding this comment

Uh oh!

Kaap10 commented Nov 3, 2026

Uh oh!

zhanyong-wan commented Nov 10, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants