Function grid_sample_3d

python/tvm/topi/testing/grid_sample_python.py:233–378

r"""grid_sample_3d for NCDHW layout

(
    data: np.ndarray,
    grid: np.ndarray,
    method="bilinear",
    layout="NCDHW",
    padding_mode="zeros",
    align_corners=True,
)

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: