chore: import upstream snapshot with attribution

python/tvm/topi/testing/adaptive_pool_python.py

@@ -0,0 +1,130 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name, unused-argument, unused-variable
+# ruff: noqa: E741, RUF005
+"""adaptive pool in python"""
+
+import numpy as np
+
+
+def _start_index(index, odim, idim):
+    return int(np.floor(index * idim / odim))
+
+
+def _end_index(index, odim, idim):
+    return int(np.ceil((index + 1) * idim / odim))
+
+
+def _pool1d(in_size, out_size, np_data, np_op):
+    out = np.zeros(out_size).astype(np_data.dtype)
+    ow = out_size[0]
+    for l in range(ow):
+        l_start = _start_index(l, ow, in_size[0])
+        l_end = _end_index(l, ow, in_size[0])
+        l_sl = slice(l_start, l_end)
+        out[l] = np_op(np_data[l_sl])
+    return out
+
+
+def _pool2d(in_size, out_size, np_data, np_op):
+    out = np.zeros(out_size).astype(np_data.dtype)
+    oh, ow = out_size
+    for k in range(oh):
+        k_start = _start_index(k, oh, in_size[0])
+        k_end = _end_index(k, oh, in_size[0])
+        k_sl = slice(k_start, k_end)
+        for l in range(ow):
+            l_start = _start_index(l, ow, in_size[1])
+            l_end = _end_index(l, ow, in_size[1])
+            l_sl = slice(l_start, l_end)
+            out[k, l] = np_op(np_data[k_sl, l_sl])
+    return out
+
+
+def _pool3d(in_size, out_size, np_data, np_op):
+    out = np.zeros(out_size).astype(np_data.dtype)
+    od, oh, ow = out_size
+    for m in range(od):
+        m_start = _start_index(m, od, in_size[0])
+        m_end = _end_index(m, od, in_size[0])
+        m_sl = slice(m_start, m_end)
+        for k in range(oh):
+            k_start = _start_index(k, oh, in_size[1])
+            k_end = _end_index(k, oh, in_size[1])
+            k_sl = slice(k_start, k_end)
+            for l in range(ow):
+                l_start = _start_index(l, ow, in_size[2])
+                l_end = _end_index(l, ow, in_size[2])
+                l_sl = slice(l_start, l_end)
+                out[m, k, l] = np_op(np_data[m_sl, k_sl, l_sl])
+    return out
+
+
+def adaptive_pool_channel_first(np_data, out_size, pool_op, np_op):
+    """The reference function for adaptive pool, channel first layout"""
+    ishape = np_data.shape
+    n, c = ishape[:2]
+    oshape = (n, c) + out_size
+    np_out = np.zeros(oshape).astype(np_data.dtype)
+
+    for i in range(n):
+        for j in range(c):
+            np_out[i, j] = pool_op(ishape[2:], out_size, np_data[i, j], np_op)
+
+    return np_out
+
+
+def adaptive_pool_channel_last(np_data, out_size, pool_op, np_op):
+    """The reference function for adaptive pool, channel last layout"""
+    ishape = np_data.shape
+    n, c = ishape[0], ishape[-1]
+    oshape = (n,) + out_size + (c,)
+    np_out = np.zeros(oshape).astype(np_data.dtype)
+
+    for i in range(n):
+        for j in range(c):
+            if len(out_size) == 1:
+                np_out[i, :, j] = pool_op(ishape[1:-1], out_size, np_data[i, :, j], np_op)
+            elif len(out_size) == 2:
+                np_out[i, :, :, j] = pool_op(ishape[1:-1], out_size, np_data[i, :, :, j], np_op)
+            else:
+                np_out[i, :, :, :, j] = pool_op(
+                    ishape[1:-1], out_size, np_data[i, :, :, :, j], np_op
+                )
+
+    return np_out
+
+
+def adaptive_pool(np_data, out_size, pool_type, layout):
+    """The reference function for adaptive pool, for 2d and 3d"""
+    if isinstance(out_size, int):
+        out_size = (out_size,)
+    if len(out_size) == 1:
+        pool_op = _pool1d
+    elif len(out_size) == 2:
+        pool_op = _pool2d
+    else:
+        assert len(out_size) == 3
+        pool_op = _pool3d
+
+    np_op = np.mean if pool_type == "avg" else np.max
+
+    if layout in ["NCW", "NCHW", "NCDHW"]:
+        return adaptive_pool_channel_first(np_data, out_size, pool_op, np_op)
+
+    assert layout in ["NWC", "NHWC", "NDHWC"]
+    return adaptive_pool_channel_last(np_data, out_size, pool_op, np_op)