chore: import upstream snapshot with attribution

python/tvm/topi/testing/grid_sample_python.py

@@ -0,0 +1,398 @@
+# 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, line-too-long, unused-variable, too-many-locals
+"""affine_grid and grid_sample operators in python"""
+
+import math
+
+import numpy as np
+
+
+def affine_grid_python(data, target_shape):
+    yv, xv = np.meshgrid(np.arange(target_shape[0]), np.arange(target_shape[1]))
+    yv = yv.T * 2 / (target_shape[0] - 1) - 1
+    xv = xv.T * 2 / (target_shape[1] - 1) - 1
+    ones = np.ones_like(xv)
+    grid = np.stack([xv, yv, ones]).reshape(3, -1)
+    return data.reshape(-1, 3).dot(grid).reshape(data.shape[0], 2, *target_shape)
+
+
+def grid_sample_2d(
+    data: np.ndarray,
+    grid: np.ndarray,
+    method="bilinear",
+    layout="NCHW",
+    padding_mode="zeros",
+    align_corners=True,
+):
+    r"""grid_sample_2d for NCHW layout"""
+
+    assert method in ("bilinear", "nearest", "bicubic"), f"{method} is not supported"
+    assert layout == "NCHW"
+    assert padding_mode in ("zeros", "border", "reflection"), f"{padding_mode} is not supported"
+    assert len(data.shape) == len(grid.shape) == 4
+
+    batch, channel = data.shape[:2]
+    in_height, in_width = data.shape[2:]
+    out_height, out_width = grid.shape[2:]
+    out_shape = [batch, channel, out_height, out_width]
+    out = np.zeros(out_shape)
+
+    def _get_pixel(b, c, h, w):
+        if 0 <= h <= in_height - 1 and 0 <= w <= in_width - 1:
+            return data[b, c, h, w]
+        return 0
+
+    def _unnormalize(h, w):
+        if align_corners:
+            new_h = (h + 1) * (in_height - 1) / 2
+            new_w = (w + 1) * (in_width - 1) / 2
+        else:
+            new_h = -0.5 + (h + 1) * in_height / 2
+            new_w = -0.5 + (w + 1) * in_width / 2
+        return (new_h, new_w)
+
+    def _clip_coordinates(x, size):
+        return min(max(x, 0), size - 1)
+
+    def _reflect_coordinates(i, size):
+        def __refelection(i, size, corner_start):
+            def __reflect(index, size, corner_start):
+                index_align_corner = abs(corner_start - index)
+                size_times = index_align_corner // size
+                even = size_times % 2 == 0
+                extra = index_align_corner - size_times * size
+                return extra + corner_start if even else size - extra + corner_start
+
+            if corner_start <= i <= size + corner_start:
+                new_i = i
+            else:
+                new_i = __reflect(i, size, corner_start)
+            return new_i
+
+        if align_corners:
+            x = __refelection(i, size - 1, 0)
+        else:
+            x = __refelection(i, size, -0.5)
+        return x
+
+    def _compute_source_index(b, h, w):
+        y = grid[b, 1, h, w]
+        x = grid[b, 0, h, w]
+        y, x = _unnormalize(y, x)
+
+        if padding_mode == "reflection":
+            y = _reflect_coordinates(y, in_height)
+            x = _reflect_coordinates(x, in_width)
+            y = _clip_coordinates(y, in_height)
+            x = _clip_coordinates(x, in_width)
+        elif padding_mode == "border":
+            y = _clip_coordinates(y, in_height)
+            x = _clip_coordinates(x, in_width)
+
+        return (y, x)
+
+    def _nearest_sample():
+        for _b in range(batch):
+            for _c in range(channel):
+                for _h in range(out_height):
+                    for _w in range(out_width):
+                        y, x = _compute_source_index(_b, _h, _w)
+                        # python round is not used here,
+                        # beacause it is done toward the even choice
+                        new_y = int(y + 0.5) if y > 0 else int(y - 0.5)
+                        new_x = int(x + 0.5) if x > 0 else int(x - 0.5)
+                        out[_b, _c, _h, _w] = _get_pixel(_b, _c, new_y, new_x)
+
+    def _bilinear_sample():
+        for _b in range(batch):
+            for _c in range(channel):
+                for _h in range(out_height):
+                    for _w in range(out_width):
+                        y, x = _compute_source_index(_b, _h, _w)
+                        y0 = math.floor(y)
+                        x0 = math.floor(x)
+                        y1 = y0 + 1
+                        x1 = x0 + 1
+
+                        out[_b, _c, _h, _w] = (
+                            _get_pixel(_b, _c, y0, x0) * (1.0 - (y - y0)) * (1.0 - (x - x0))
+                            + _get_pixel(_b, _c, y0, x1) * (1.0 - (y - y0)) * (x - x0)
+                            + _get_pixel(_b, _c, y1, x0) * (y - y0) * (1.0 - (x - x0))
+                            + _get_pixel(_b, _c, y1, x1) * (y - y0) * (x - x0)
+                        )
+
+    def _bicubic_sample():
+        A = -0.75
+
+        def cubic_weight_1(x_fraction):
+            return ((A + 2) * x_fraction - (A + 3)) * x_fraction * x_fraction + 1
+
+        def cubic_weight_2(x_fraction):
+            return ((A * x_fraction - 5 * A) * x_fraction + 8 * A) * x_fraction - 4 * A
+
+        def cubic_interp_1d(pixel_0, pixel_1, pixel_2, pixel_3, x_fraction):
+            weights = [0] * 4
+            weights[0] = cubic_weight_2(x_fraction + 1)
+            weights[1] = cubic_weight_1(x_fraction)
+            weights[2] = cubic_weight_1(1 - x_fraction)
+            weights[3] = cubic_weight_2(2 - x_fraction)
+
+            return (
+                pixel_0 * weights[0]
+                + pixel_1 * weights[1]
+                + pixel_2 * weights[2]
+                + pixel_3 * weights[3]
+            )
+
+        def coefficients_along_x(x_floor, y_floor, x_fraction):
+            coefficients = [0] * 4
+
+            for i in range(4):
+                y_ = y_floor - 1 + i
+                x_0 = x_floor - 1
+                x_1 = x_floor + 0
+                x_2 = x_floor + 1
+                x_3 = x_floor + 2
+
+                if padding_mode == "border":
+                    y_ = _clip_coordinates(y_, in_height)
+                    x_0 = _clip_coordinates(x_0, in_width)
+                    x_1 = _clip_coordinates(x_1, in_width)
+                    x_2 = _clip_coordinates(x_2, in_width)
+                    x_3 = _clip_coordinates(x_3, in_width)
+
+                elif padding_mode == "reflection":
+                    y_ = _reflect_coordinates(y_, in_height)
+                    x_0 = _reflect_coordinates(x_0, in_width)
+                    x_1 = _reflect_coordinates(x_1, in_width)
+                    x_2 = _reflect_coordinates(x_2, in_width)
+                    x_3 = _reflect_coordinates(x_3, in_width)
+
+                    y_ = int(_clip_coordinates(y_, in_height))
+                    x_0 = int(_clip_coordinates(x_0, in_width))
+                    x_1 = int(_clip_coordinates(x_1, in_width))
+                    x_2 = int(_clip_coordinates(x_2, in_width))
+                    x_3 = int(_clip_coordinates(x_3, in_width))
+
+                coefficients[i] = cubic_interp_1d(
+                    _get_pixel(_b, _c, y_, x_0),
+                    _get_pixel(_b, _c, y_, x_1),
+                    _get_pixel(_b, _c, y_, x_2),
+                    _get_pixel(_b, _c, y_, x_3),
+                    x_fraction,
+                )
+            return coefficients
+
+        for _b in range(batch):
+            for _c in range(channel):
+                for _h in range(out_height):
+                    for _w in range(out_width):
+                        y = grid[_b, 1, _h, _w]
+                        x = grid[_b, 0, _h, _w]
+                        y, x = _unnormalize(y, x)
+                        y_floor = math.floor(y)
+                        x_floor = math.floor(x)
+                        y_fraction = y - y_floor
+                        x_fraction = x - x_floor
+
+                        coefficients = coefficients_along_x(x_floor, y_floor, x_fraction)
+
+                        out[_b, _c, _h, _w] = cubic_interp_1d(
+                            coefficients[0],
+                            coefficients[1],
+                            coefficients[2],
+                            coefficients[3],
+                            y_fraction,
+                        )
+
+    if method == "bilinear":
+        _bilinear_sample()
+    elif method == "nearest":
+        _nearest_sample()
+    else:  # mode == "bicubic":
+        _bicubic_sample()
+
+    return out
+
+
+def grid_sample_3d(
+    data: np.ndarray,
+    grid: np.ndarray,
+    method="bilinear",
+    layout="NCDHW",
+    padding_mode="zeros",
+    align_corners=True,
+):
+    r"""grid_sample_3d for NCDHW layout"""
+
+    assert method in ("bilinear", "nearest"), f"{method} is not supported"
+    assert layout == "NCDHW"
+    assert padding_mode in ("zeros", "border", "reflection"), f"{padding_mode} is not supported"
+    assert len(data.shape) == len(grid.shape) == 5
+
+    batch, channel = data.shape[:2]
+    in_depth, in_height, in_width = data.shape[2:]
+    out_depth, out_height, out_width = grid.shape[2:]
+    out_shape = [batch, channel, out_depth, out_height, out_width]
+    out = np.zeros(out_shape)
+
+    def _get_pixel(b, c, d, h, w):
+        if 0 <= d <= in_depth - 1 and 0 <= h <= in_height - 1 and 0 <= w <= in_width - 1:
+            return data[b, c, d, h, w]
+        return 0
+
+    def _unnormalize(d, h, w):
+        if align_corners:
+            new_d = (d + 1) * (in_depth - 1) / 2
+            new_h = (h + 1) * (in_height - 1) / 2
+            new_w = (w + 1) * (in_width - 1) / 2
+        else:
+            new_d = -0.5 + (d + 1) * in_depth / 2
+            new_h = -0.5 + (h + 1) * in_height / 2
+            new_w = -0.5 + (w + 1) * in_width / 2
+        return (new_d, new_h, new_w)
+
+    def _clip_coordinates(x, size):
+        return min(max(x, 0), size - 1)
+
+    def _reflect_coordinates(i, size):
+        def __refelection(i, size, corner_start):
+            def __reflect(index, size, corner_start):
+                index_align_corner = abs(corner_start - index)
+                size_times = index_align_corner // size
+                even = size_times % 2 == 0
+                extra = index_align_corner - size_times * size
+                return extra + corner_start if even else size - extra + corner_start
+
+            if corner_start <= i <= size + corner_start:
+                new_i = i
+            else:
+                new_i = __reflect(i, size, corner_start)
+            return new_i
+
+        if align_corners:
+            x = __refelection(i, size - 1, 0)
+        else:
+            x = __refelection(i, size, -0.5)
+        return x
+
+    def _compute_source_index(b, d, h, w):
+        z = grid[b, 2, d, h, w]
+        y = grid[b, 1, d, h, w]
+        x = grid[b, 0, d, h, w]
+        z, y, x = _unnormalize(z, y, x)
+
+        if padding_mode == "reflection":
+            z = _reflect_coordinates(z, in_depth)
+            y = _reflect_coordinates(y, in_height)
+            x = _reflect_coordinates(x, in_width)
+            z = _clip_coordinates(z, in_depth)
+            y = _clip_coordinates(y, in_height)
+            x = _clip_coordinates(x, in_width)
+        elif padding_mode == "border":
+            z = _clip_coordinates(z, in_depth)
+            y = _clip_coordinates(y, in_height)
+            x = _clip_coordinates(x, in_width)
+        return (z, y, x)
+
+    def _nearest_sample():
+        for _b in range(batch):
+            for _c in range(channel):
+                for _d in range(out_depth):
+                    for _h in range(out_height):
+                        for _w in range(out_width):
+                            z, y, x = _compute_source_index(_b, _d, _h, _w)
+                            # python round is not used here,
+                            # beacause it is done toward the even choice
+                            new_z = int(z + 0.5) if z > 0 else int(z - 0.5)
+                            new_y = int(y + 0.5) if y > 0 else int(y - 0.5)
+                            new_x = int(x + 0.5) if x > 0 else int(x - 0.5)
+                            out[_b, _c, _d, _h, _w] = _get_pixel(_b, _c, new_z, new_y, new_x)
+
+    def _triilinear_sample():
+        for _b in range(batch):
+            for _c in range(channel):
+                for _d in range(out_depth):
+                    for _h in range(out_height):
+                        for _w in range(out_width):
+                            z, y, x = _compute_source_index(_b, _d, _h, _w)
+                            z0 = math.floor(z)
+                            y0 = math.floor(y)
+                            x0 = math.floor(x)
+                            z1 = z0 + 1
+                            y1 = y0 + 1
+                            x1 = x0 + 1
+
+                            out[_b, _c, _d, _h, _w] = (
+                                _get_pixel(_b, _c, z0, y0, x0)
+                                * (1 - (x - x0))
+                                * (1 - (y - y0))
+                                * (1 - (z - z0))
+                                + _get_pixel(_b, _c, z0, y0, x1)
+                                * (x - x0)
+                                * (1 - (y - y0))
+                                * (1 - (z - z0))
+                                + _get_pixel(_b, _c, z1, y1, x0)
+                                * (1 - (x - x0))
+                                * (y - y0)
+                                * (z - z0)
+                                + _get_pixel(_b, _c, z1, y1, x1) * (x - x0) * (y - y0) * (z - z0)
+                                + _get_pixel(_b, _c, z0, y1, x0)
+                                * (1 - (x - x0))
+                                * (y - y0)
+                                * (1 - (z - z0))
+                                + _get_pixel(_b, _c, z1, y0, x1)
+                                * (x - x0)
+                                * (1 - (y - y0))
+                                * (z - z0)
+                                + _get_pixel(_b, _c, z1, y0, x0)
+                                * (1 - (x - x0))
+                                * (1 - (y - y0))
+                                * (z - z0)
+                                + _get_pixel(_b, _c, z0, y1, x1)
+                                * (x - x0)
+                                * (y - y0)
+                                * (1 - (z - z0))
+                            )
+
+    if method == "bilinear":
+        _triilinear_sample()
+    else:  # method == "nearest":
+        _nearest_sample()
+
+    return out
+
+
+def grid_sample_python(
+    data: np.ndarray,
+    grid: np.ndarray,
+    method="bilinear",
+    layout="NCHW",
+    padding_mode="zeros",
+    align_corners=True,
+):
+    r"""grid_sample_3d for NCDHW layout or grid_sample_2d for NCHW layout"""
+
+    if len(data.shape) == 4:
+        grid_sample = grid_sample_2d
+    elif len(data.shape) == 5:
+        grid_sample = grid_sample_3d
+    else:
+        raise ValueError("invalid shape")
+
+    return grid_sample(data, grid, method, layout, padding_mode, align_corners)