Function poolnd_python

python/tvm/topi/testing/poolnd_python.py:126–209

Ground truth pooling operator impelmented in numpy.

(
    np_data: np.array,
    kernel: tuple[int],
    strides: tuple[int],
    dilation: tuple[int],
    padding_before: tuple[int],
    padding_after: tuple[int],
    pool_type: str,
    count_include_pad: bool = True,
    ceil_mode: bool = False,
    dtype: str = "float32",
    layout: str | None = None,
)

def poolnd_python(
    np_data: np.array,
    kernel: tuple[int],
    strides: tuple[int],
    dilation: tuple[int],
    padding_before: tuple[int],
    padding_after: tuple[int],
    pool_type: str,
    count_include_pad: bool = True,
    ceil_mode: bool = False,
    dtype: str = "float32",
    layout: str | None = None,
) -> np.array:
    """Ground truth pooling operator impelmented in numpy."""

    np_data = _convert_from_layout(np_data, layout)

    out_shape = [np_data.shape[0], np_data.shape[1]]
    for dim in range(2, len(np_data.shape)):
        i = dim - 2
        val = (
            float(
                np_data.shape[dim]
                - (kernel[i] - 1) * dilation[i]
                - 1
                + padding_before[i]
                + padding_after[i]
            )
            / strides[i]
        )

        if ceil_mode:
            out_shape.append(int(math.ceil(val) + 1))
        else:
            out_shape.append(int(math.floor(val) + 1))
    out_shape = tuple(out_shape)

    # Create a padded array, and a boolean mask showing which values are padded values
    pad_value = 0
    if pool_type == "max" and not count_include_pad:
        pad_value = tvm.te.min_value(dtype).value
    pad_data = pad_tensor(np_data, pad_value, padding_before, padding_after, dtype)
    pad_map = pad_tensor(np.ones_like(np_data), 0, padding_before, padding_after, "bool")

    # Create iterator which gives all indices for output array
    dim_iterators = []
    for spatial_dimension in range(2, len(np_data.shape)):
        dim_iterators.append(range(out_shape[spatial_dimension]))
    coord_iterator = itertools.product(*dim_iterators)

    ret_np = np.zeros(shape=out_shape).astype(dtype)
    for coordinate in coord_iterator:
        # Get index into the values that any pool operation will use for given coordinate
        np_index = get_slice(
            spatial_dimensions=len(out_shape) - 2,
            pad_np=pad_data,
            dim_coord=coordinate,
            kernel=kernel,