Function pad_tensor

python/tvm/topi/testing/poolnd_python.py:102–123

Pad the spatial dimensions of the given array.

(
    np_arr: np.array,
    pad_value: float,
    padding_before: list[int],
    padding_after: list[int],
    dtype: str,
)

def pad_tensor(
    np_arr: np.array,
    pad_value: float,
    padding_before: list[int],
    padding_after: list[int],
    dtype: str,
) -> np.array:
    """Pad the spatial dimensions of the given array."""
    orig_shape = list(np_arr.shape)
    padded_shape = list(np_arr.shape)
    n = len(orig_shape)
    for dim in range(2, n):
        i = dim - 2
        padded_shape[dim] += padding_after[i] + padding_before[i]

    pad_np = (np.zeros(shape=padded_shape) + pad_value).astype(dtype)
    ranges_it = [range(padded_shape[0]), range(padded_shape[1])]
    for dim in range(2, n):
        i = dim - 2
        ranges_it.append(range(padding_before[i], padding_before[i] + orig_shape[dim]))
    pad_np[np.ix_(*ranges_it)] = np_arr
    return pad_np


def poolnd_python(