Repository: onnx/onnx
Pull request: Add check on dimensions in Gemm opset 6
Fix https://github.com/onnx/onnx/issues/6185




Problem statement:
Shape Inference crash on Gemm
# Bug Report

### Describe the bug
```
import onnx
import onnx.parser

model = """
<
   ir_version: 10,
   opset_import: ["" : 6]
>
graph (double[2,1] in0, double in1, double[2] in2) => () {
   out0 = Gemm <alpha: float = 1, beta: float = -693.752, broadcast: int = -436, transB: int = 823> (in0, in1, in2)
}
"""

onnx.shape_inference.infer_shapes(onnx.parser.parse_model(model))
```
crashes with a segmentation fault.

### System information
- OS Platform and Distribution (*. Linux Ubuntu 20.04*): Linux Ubuntu 20.04
- ONNX version (*e.g. 1.13*): 1.16.1
- Python version: 3.10.12

### Expected behavior
No crash, but an error that the model is invalid

Hints from issue discussion:

Commit message: Add check on dimensions in Gemm opset 6

Signed-off-by: Ganesan Ramalingam <grama@microsoft.com>

Patch to review:
diff --git a/onnx/defs/math/old.cc b/onnx/defs/math/old.cc
index d63070f9c04..0e428c56b82 100644
--- a/onnx/defs/math/old.cc
+++ b/onnx/defs/math/old.cc
@@ -3304,10 +3304,19 @@ ONNX_OPERATOR_SET_SCHEMA(
             auto transBAttr = ctx.getAttribute("transB");
             bool transB = transBAttr ? static_cast<int>(transBAttr->i()) != 0 : false;
 
+            auto& first_input_shape = getInputShape(ctx, 0);
+            auto& second_input_shape = getInputShape(ctx, 1);
+            if (first_input_shape.dim_size() != 2) {
+              fail_shape_inference("First input does not have rank 2");
+            }
+            if (second_input_shape.dim_size() != 2) {
+              fail_shape_inference("Second input does not have rank 2");
+            }
+
             *ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape()->add_dim() =
-                ctx.getInputType(0)->tensor_type().shape().dim(transA ? 1 : 0);
+                first_input_shape.dim(transA ? 1 : 0);
             *ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape()->add_dim() =
-                ctx.getInputType(1)->tensor_type().shape().dim(transB ? 0 : 1);
+                second_input_shape.dim(transB ? 0 : 1);
           } else if (
               hasInputShape(ctx, 2) &&
               (!ctx.getAttribute("broadcast") || static_cast<int>(ctx.getAttribute("broadcast")->i()) == 0)) {


Code review comments:
diff_hunk: @@ -3304,10 +3304,19 @@ ONNX_OPERATOR_SET_SCHEMA(
             auto transBAttr = ctx.getAttribute("transB");
             bool transB = transBAttr ? static_cast<int>(transBAttr->i()) != 0 : false;
 
+            auto& first_input_shape = getInputShape(ctx, 0);
+            auto& second_input_shape = getInputShape(ctx, 1);
+            if (first_input_shape.dim_size() != 2) {
+              fail_shape_inference("First input does not have rank 2");
+            }
+            if (second_input_shape.dim_size() != 2) {
+              fail_shape_inference("Second input does not have rank 2");
path: onnx/defs/math/old.cc
text: @user1:
I would include the actual shape in the error message to help diagnose issues.

@author:
Done, by calling utility which does that.