Skip to content

Validate LayerNormalization axis bounds - #32555

Open
Akshay Sonawane (apsonawane) wants to merge 1 commit into
mainfrom
fix/layernorm-axis-shape-inference
Open

Validate LayerNormalization axis bounds#32555
Akshay Sonawane (apsonawane) wants to merge 1 commit into
mainfrom
fix/layernorm-axis-shape-inference

Conversation

@apsonawane

Copy link
Copy Markdown
Contributor

This pull request addresses validation and testing improvements for the LayerNormalization operator, specifically handling invalid axis values. The main change ensures that the operator correctly detects and reports an error when the axis attribute exceeds the input tensor's rank. Additionally, a new unit test is added to verify this behavior.

Validation logic update:

  • Updated the axis validation in onnx/defs/nn/defs.cc to check if axis is less than 0 or greater than or equal to the input tensor rank, and to fail shape inference with an appropriate error message if so.

Unit testing:

  • Added a new test case LayerNorm_AxisExceedsRank in onnxruntime/test/contrib_ops/layer_norm_op_test.cc to ensure that providing an out-of-bounds axis value results in a failure with the expected error message.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

The test does not cover the exact axis == rank boundary introduced by the new condition.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds LayerNormalization shape-inference validation for axes at or above the input rank.

Changes:

  • Updates both ONNX dependency patches with the bounds check.
  • Adds an invalid-axis unit test.
File summaries
File Description
onnxruntime/test/contrib_ops/layer_norm_op_test.cc Tests rejection of an oversized axis.
cmake/vcpkg-ports/onnx/binskim.patch Applies validation to the vcpkg ONNX source.
cmake/patches/onnx/onnx.patch Applies validation to the standard ONNX dependency.
Review details
  • Files reviewed: 3/3 changed files
  • Comments generated: 1
  • Review effort level: Balanced

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment on lines +858 to +867
TEST(LayerNormTest, LayerNorm_AxisExceedsRank) {
OpTester test("LayerNormalization", 17);
test.AddAttribute<int64_t>("axis", 0xFFFFFFFFLL);
test.AddInput<float>("X", {1, 2}, {1.0f, 2.0f});
test.AddInput<float>("Scale", {2}, {1.0f, 1.0f});
test.AddOutput<float>("Y", {1, 2}, {0.0f, 0.0f});
test.AddOutput<float>("Mean", {1, 1}, {0.0f});

test.Run(OpTester::ExpectResult::kExpectFailure, "Unexpected axis value");
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants