From 89c6da99b5d229d67b465a2ecf158a4254932cd3 Mon Sep 17 00:00:00 2001
From: Ashwini Khade <askhade@microsoft.com>
Date: Tue, 30 Jun 2020 08:21:20 -0700
Subject: [PATCH] fix output shape calc for matmul (#4362)

---
 onnxruntime/core/providers/cpu/math/matmul_helper.h | 2 +-
 onnxruntime/test/providers/cpu/math/matmul_test.cc  | 7 +++++++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/onnxruntime/core/providers/cpu/math/matmul_helper.h b/onnxruntime/core/providers/cpu/math/matmul_helper.h
index fca73c2d99..2d8d6c14aa 100644
--- a/onnxruntime/core/providers/cpu/math/matmul_helper.h
+++ b/onnxruntime/core/providers/cpu/math/matmul_helper.h
@@ -29,7 +29,7 @@ class MatMulComputeHelper {
     // A: [M1, M2, ... K], B: [N, K]^T
     // A: [M1, M2, ... K], B: [1, ..., 1, K, N]
     // A: [M1, M2, ... K], B: [1, ..., 1, N, K]^T
-    if (!transa && left_num_dims >= 2 && right_num_dims >= 2 &&
+    if (!transa && left_num_dims >= 2 && right_num_dims >= 2 && left_num_dims >= right_num_dims &&
         right_shape.SizeToDimension(right_num_dims - 1) == right_shape[right_num_dims - 2]) {
       M_ = left_shape.SizeToDimension(left_num_dims - 1);
       K_ = left_shape[left_num_dims - 1];
diff --git a/onnxruntime/test/providers/cpu/math/matmul_test.cc b/onnxruntime/test/providers/cpu/math/matmul_test.cc
index 1e6bece98c..418dee24f0 100644
--- a/onnxruntime/test/providers/cpu/math/matmul_test.cc
+++ b/onnxruntime/test/providers/cpu/math/matmul_test.cc
@@ -77,6 +77,13 @@ std::vector<MatMulTestData<T>> GenerateTestCases()
     {2, 2, 4},
     {20, 23, 26, 29, 56, 68, 80, 92, 92, 113, 134, 155, 128, 158, 188, 218}});
 
+  test_cases.push_back(
+    {"test 2D special 3",
+    {2, 6},
+    {1, 1, 6, 1},
+    {1, 1, 2, 1},
+    {55, 145}});
+
   return test_cases;
 }