Max pooling using numpy

Question

I am trying to implement max pooling using numpy and below is my snippet.

However I am getting an error cannot reshape array of size 2883 into shape (2,3,31,31)

My input is 2X3X32X32.

my method call is maxpool(x, kernel_size=2, stride=1, padding=0)

def maxpool(x, kernel_size, stride, padding):
"""
Args:
    x: torch tensor with size (N, C_in, H_in, W_in),
    kernel_size: size of the window to take a max over, 
    stride: stride of the window,
    padding: implicit zero padding to be added on both sides,
    
Return:
    y: torch tensor of size (N, C_out, H_out, W_out).
"""
y = None
x_height = x.shape[2]
x_width = x.shape[3]
channels = x.shape[1]
n = x.shape[0]
y_height = int((x_height-kernel_size)/stride)+1
y_width = int((x_width-kernel_size)/stride)+1
output = np.zeros((channels,y_height*y_width))
c = 0
for height in range(0, x_height, stride):
  if height + kernel_size <= x_height:
    image_rectangle = x[0, :, height:height + kernel_size, :]
    for width in range(0, x_width, stride):
      if width+kernel_size <= x_width:
        image_square = image_rectangle[:, :, width:width + kernel_size]
        image_flatten = image_square.reshape(-1, 1)
        output[:, c:c+1] = np.array([float(max(i)) for i in np.split(image_flatten, channels)]).reshape(-1, 1)
        c+=1
y = np.array(np.hsplit(output, 1)).reshape((n, channels, y_height, y_width))
return y