How can I vectorize overlapping window gradient accumulation in NumPy without using for loops?

I'm implementing a custom 2D pooling layer (with channels) in Python using only NumPy. In the backward pass, I use an as_strided view to extract overlapping windows from a padded input, and I need to add each window's gradient back to its original location. However, since the windows may overlap (depending on stride, padding, and kernel size), the current Python for‑loop approach causes incorrect gradient accumulation. I want to remove the loop and fully vectorize this operation.

Below is the class implementation I developed, which includes both the forward and backward passes:

PYTHONCopied
import numpy as np

def maxPool(partition: np.ndarray, backwards: bool = False):
        max_val = np.max(partition)
        if backwards:
            mask = np.zeros(partition.shape, dtype=partition.dtype)
            mask[partition.argmax()] = 1
            return mask
        return max_val

class Pooling2D():
    def __init__(self, pool_funct: callable, kernel_size: tuple[int, int], 
                 input_size: tuple[int, int, int], strides: tuple[int, int] = (1, 1), padding: int = 0) -> None:
        self.funct = pool_funct
        self.in_size = input_size
        self.out_size = (input_size[0],
                        (input_size[1] - kernel_size[0] + 2*padding) // strides[0] + 1, 
                        (input_size[2] - kernel_size[1] + 2*padding) // strides[1] + 1)
        self.kernel_height, self.kernel_width = kernel_size
        self.window_shape = (*self.out_size, 
                             self.kernel_height, 
                             self.kernel_width)
        self.stride_h, self.stride_w = strides
        self.padding = padding

    def forward(self, x: np.ndarray) -> np.ndarray:
        x_padded = np.pad(x, pad_width=((0, 0), (self.padding, self.padding), (self.padding, self.padding)), mode='constant')
        strides = (x_padded.strides[0], 
                   self.stride_h * x_padded.strides[1], 
                   self.stride_w * x_padded.strides[2], 
                   x_padded.strides[1], 
                   x_padded.strides[2])

        # Split into windows and apply the pooling function to each window.
        A_s = np.lib.stride_tricks.as_strided(x_padded, shape=self.window_shape, strides=strides)
        pool_val =  np.apply_along_axis(self.funct, -1, A_s.reshape(*self.out_size, -1))

        self.last_in = x
        self.last_out = pool_val
        return pool_val

    def backward(self, dx: np.ndarray) -> np.ndarray:
        x_padded = np.pad(self.last_in, ((0, 0), (self.padding, self.padding), (self.padding, self.padding)), mode="constant")
        strides = (x_padded.strides[0], 
                   self.stride_h * x_padded.strides[1], 
                   self.stride_w * x_padded.strides[2], 
                   x_padded.strides[1], 
                   x_padded.strides[2])

        # Window frames for previous input / Mask creation
        main_windows = np.lib.stride_tricks.as_strided(x_padded, self.window_shape, strides, writeable=False)
        mask = np.apply_along_axis(self.funct, -1, main_windows.reshape(*self.out_size, -1), backwards=True).reshape(main_windows.shape)

        # Use mask to distribute the gradient into the mask, reshaped into (channels, kernel height, kernel width, num of windows)
        pre_grad = np.einsum("ghw,ghwxy->ghwxy", dx, mask).transpose(0, 3, 4, 1, 2)
        pre_grad = pre_grad.reshape(*pre_grad.shape[:3], -1) 

        # Zero array of original size (channels, in height, in width)
        final_grad = np.zeros_like(x_padded)

        # _____________________________________________
        # - - - - - THE PART TO BE VECTORIZED - - - - -
        # _____________________________________________
        right = 0
        down = 0
        for i in range(pre_grad.shape[3]):
            if right+self.kernel_width > x_padded.shape[2]:
                right = 0
                down += self.stride_h
            final_grad[:, down:down+self.kernel_height, 
                          right:right+self.kernel_width] += pre_grad[:, :, :, i]
            right += self.stride_w

        return final_grad[:, self.padding:-self.padding, self.padding:-self.padding] if self.padding > 0 else final_grad

I also attempted a second approach using another as_strided call with reconstructed strides from the original array. For example:

PYTHONCopied
# Get the strides from the original array
_, Hs, Ws = x.strides
new_strides = (Hs*Ws*pre_grad.itemsize, Ws*pre_grad.itemsize, pre_grad.itemsize)

# Do a secondary as_strided
final_grad = np.lib.stride_tricks.as_strided(pre_grad, x.shape, new_strides)

However, with this method, when using a (1, 1) stride and a (3, 3) kernel, overlapping regions are not summed correctly and values end up in the wrong locations. My question is: can I vectorize the gradient accumulation (the part marked above) to correctly handle overlaps without resorting to a Python loop?