update pool tests

titaiwangms · Jan 22, 2025 · f97ca29 · f97ca29
1 parent 505a942
commit f97ca29
Show file tree

Hide file tree

Showing 10 changed files with 24 additions and 17 deletions.
diff --git a/...l_3d_dilations_large_count_include_pad_is_1_ceil_mode_is_True/test_data_set_0/output_0.pb b/...l_3d_dilations_large_count_include_pad_is_1_ceil_mode_is_True/test_data_set_0/output_0.pb
diff --git a/onnx/backend/test/data/node/test_image_decoder_decode_jpeg2k_rgb/test_data_set_0/input_0.pb b/onnx/backend/test/data/node/test_image_decoder_decode_jpeg2k_rgb/test_data_set_0/input_0.pb
diff --git a/onnx/backend/test/data/node/test_lppool_1d_default/test_data_set_0/output_0.pb b/onnx/backend/test/data/node/test_lppool_1d_default/test_data_set_0/output_0.pb
@@ -1,2 +1,2 @@
-ByJ����?��?aK@��&@���?�q�?̌s?;�)>wU�>�5�>5�?;��?8C?���>���>�?3h�?B�>8.^?�g%@�L$@&sy?TK�?�@i@')�?�@>�P�?k/�?�&�?���>|w@���?V]�>���?/�?˙�?���>GG�?Dv�?/�?X@�,�?�?���?�;�?��?���?qBf?�>k?�z?]�?"�?�S�>Q��>�i�>(?��+?�4?I�4?��U?�h�?~t�>���?�>�?rr?mHh?p�:?��:?��? ��?ܟ?�p:?�h~?��r?��?hӟ>%�?�E�?w�p?�r�?��?Cd
-@�^@��?�E�?6`�?tp�?|R�?��?E�z?2~?,O<?
+ByJ����?��?aK@��&@���?�q�?͌s?;�)>wU�>�5�>5�?;��?8C?���>���>�?3h�?B�>7.^?�g%@�L$@%sy?TK�?�@i@')�?�@>�P�?k/�?�&�?���>|w@���?V]�>���?/�?˙�?���>GG�?Dv�?/�?X@�,�?�?���?�;�?��?���?qBf?�>k?�z?]�?"�?�S�>P��>�i�>(?��+?�4?I�4?��U?�h�?~t�>���?�>�?rr?nHh?p�:?��:?��? ��?ܟ?�p:?�h~?��r?��?hӟ>%�?�E�?w�p?�r�?��?Cd
+@�^@��?�E�?5`�?tp�?|R�?��?E�z?2~?,O<?

diff --git a/onnx/backend/test/data/node/test_lppool_2d_default/test_data_set_0/output_0.pb b/onnx/backend/test/data/node/test_lppool_2d_default/test_data_set_0/output_0.pb
diff --git a/onnx/backend/test/data/node/test_lppool_2d_pads/test_data_set_0/output_0.pb b/onnx/backend/test/data/node/test_lppool_2d_pads/test_data_set_0/output_0.pb
diff --git a/onnx/backend/test/data/node/test_lppool_2d_same_lower/test_data_set_0/output_0.pb b/onnx/backend/test/data/node/test_lppool_2d_same_lower/test_data_set_0/output_0.pb
diff --git a/onnx/backend/test/data/node/test_lppool_2d_same_upper/test_data_set_0/output_0.pb b/onnx/backend/test/data/node/test_lppool_2d_same_upper/test_data_set_0/output_0.pb
diff --git a/onnx/backend/test/data/node/test_lppool_2d_strides/test_data_set_0/output_0.pb b/onnx/backend/test/data/node/test_lppool_2d_strides/test_data_set_0/output_0.pb
diff --git a/onnx/backend/test/data/node/test_lppool_3d_default/test_data_set_0/output_0.pb b/onnx/backend/test/data/node/test_lppool_3d_default/test_data_set_0/output_0.pb
diff --git a/onnx/reference/ops/op_pool_common.py b/onnx/reference/ops/op_pool_common.py
@@ -218,22 +218,29 @@ def pool(
         ],
     ):
         window = padded[shape[0], shape[1]]
-        elements = []
-        for i in range(spatial_size):
-            # NOTE: The if condition is to avoid the case where the window is out of bound
-            # we need to avoid the pixels that are out of bound being included in the window
-            elements.extend(
-                num
-                for num in range(
-                    strides[i] * shape[i + 2],
-                    strides[i] * shape[i + 2] + (1 + (kernel[i] - 1) * dilations[i]),
-                    dilations[i],
+        window_vals = np.array(
+            [
+                window[i]
+                for i in list(
+                    itertools.product(
+                        *[
+                            [
+                                pixel
+                                for pixel in range(
+                                    strides[i] * shape[i + 2],
+                                    strides[i] * shape[i + 2]
+                                    + (1 + (kernel[i] - 1) * dilations[i]),
+                                    dilations[i],
+                                )
+                                if pixel
+                                < x_shape[i + 2] + pads[i] + pads[spatial_size + i]
+                            ]
+                            for i in range(spatial_size)
+                        ]
+                    )
                 )
-                if num < x_shape[i + 2] + pads[i] + pads[i + spatial_size]
-            )
-            window_vals = np.array(
-                [window[indices] for indices in itertools.product(elements)]
-            )
+            ]
+        )
 
         if pooling_type == "AVG":
             f = np.average