Function maxpooling_grad

imperative/python/megengine/xla/rules/nn.py:501–544

(
    x,
    dy,
    kernel,
    stride,
    padding,
    base_dilation=None,
    kernel_dilation=None,
    expand_padding=True,
)

def maxpooling_grad(
    x,
    dy,
    kernel,
    stride,
    padding,
    base_dilation=None,
    kernel_dilation=None,
    expand_padding=True,
):
    assert base_dilation is None and kernel_dilation is None
    assert expand_padding == True
    padding = [(p, p) if isinstance(p, int) else p for p in padding]
    dxdtype, dxshape = x.dtype, x.shape
    assert dxdtype == "float32" or dxdtype == "float16"

    org_padding, new_padding = padding, padding
    if expand_padding:
        pads = [(lo, hi, 0) for (lo, hi) in padding]
        padded_x = pad(x, _get_max_identity(dxdtype), pads)
        new_padding = [(0, 0) for _ in padding]

    selector = lambda x, y: x >= y
    scatter = lambda x, y: x + y
    out = _select_and_scatter(
        padded_x,
        dy,
        np.array(0.0, dtype=dy.dtype),
        kernel,
        stride,
        new_padding,
        selector,
        scatter,
    )

    if expand_padding:
        start_indices = [lo for (lo, hi) in org_padding]
        stop_indices = [lo + d for ((lo, hi), d) in zip(org_padding, dxshape)]
        slices = [
            slice(start, stop, 1) for start, stop in zip(start_indices, stop_indices)
        ]
        out = index_with_slices(out, slices)

    return out


def avgpooling_grad(