From 9fda95fec9925445d85d79ffab54c9b35542e029 Mon Sep 17 00:00:00 2001
From: Chi Lo <54722500+chilo-ms@users.noreply.github.com>
Date: Sat, 25 Sep 2021 10:13:54 -0700
Subject: [PATCH] Change keepdims of ReduceMax/ReduceMin to always 1 when using
 quatization calibration MinMax approach (#9167)

* Change keepdims to always 1

* fix typo

* Refine code
---
 .../python/tools/quantization/calibrate.py    | 26 +++++++------------
 1 file changed, 10 insertions(+), 16 deletions(-)

diff --git a/onnxruntime/python/tools/quantization/calibrate.py b/onnxruntime/python/tools/quantization/calibrate.py
index f38e9865ad..4c147ebcfd 100644
--- a/onnxruntime/python/tools/quantization/calibrate.py
+++ b/onnxruntime/python/tools/quantization/calibrate.py
@@ -164,23 +164,17 @@ class MinMaxCalibrater(CalibraterBase):
 
         for tensor in tensors:
 
-            # When doing ReduceMax/ReduceMin, keep dimension if tensor contains dim with value of 0,
-            # for example:
-            #     dim = [ dim_value: 0 ] 
-            #  
-            # otherwise, don't keep dimension. 
-            #
+            # When doing ReduceMax/ReduceMin, ORT can't reduce on dim with value of 0 if 'keepdims' is false.
+            # To make the code simple, we always let keepdims to be 1.
+            keepdims = 1
+
+            # dim could be:
+            #   [dim_param: "batch_size", dim_value: 256, dim_value: 36, dim_value: 64],
+            #   [dim_value: 0],
+            #   ...
+            # Please see the definition of TensorShapeProto https://github.com/onnx/onnx/blob/master/onnx/onnx.proto#L651
             dim = value_infos[tensor].type.tensor_type.shape.dim
-            keepdims = 0
-            shape = ()
-            for d in dim:
-                # A dimension can be either an integer value or a symbolic variable.
-                # Dimension with integer value and value of 0 is what we are looking for to keep dimension. 
-                # Please see the def of TensorShapeProto https://github.com/onnx/onnx/blob/master/onnx/onnx.proto#L630
-                if d.WhichOneof('value') == 'dim_value' and d.dim_value == 0:
-                    keepdims = 1
-                    shape = (1,) if len(dim) == 1 else list(1 for i in range(len(dim)))
-                    break
+            shape = (1,) if len(dim) == 1 else tuple(1 for i in range(len(dim)))
 
             # Adding ReduceMin nodes
             reduce_min_name = tensor + '_ReduceMin'