How do I create a binaryMask that displays all the values which contributed to the output

Question

I've been trying to implement a MaxPoolingLayer in numpy. The problem I'm having is that I cannot create an array that contains all the indices of the max Values(which I obviously need for the backward pass). My Idea is to create said array as some sort of binaryMask which i can apply to the dInputArray to get an Array which is set to 0 on all values that did not contribute to the output of the forward pass but set to one to all values which contributed to the output of the forward pass. My Code is here:

import numpy as np

#Quellen https://numpy.org/doc/stable/reference/generated/numpy.mgrid.html
#https://lanstonchu.wordpress.com/2018/09/01/convolutional-neural-network-cnn-backward-propagation-of-the-pooling-layers/

class MaxPooling():
    def __init__(self, *, maxPoolingSize=2):
        self.maxPoolingSize = maxPoolingSize
    
    def forward(self, image):
        self._input = image.shape
        b,h, w  = image.shape
        output = np.zeros((b, h // self.maxPoolingSize, w // self.maxPoolingSize))
        self._maxOutputIndeces = np.copy(output)
        for batch in range(b):
            for y in range(0, h - h % self.maxPoolingSize, self.maxPoolingSize):
                for x in range(0, w - w % self.maxPoolingSize, self.maxPoolingSize):
                    Y, X = y // self.maxPoolingSize, x // self.maxPoolingSize
                    output[batch, Y, X] = np.max(image[batch, y:y+self.maxPoolingSize, x:x+self.maxPoolingSize])
                    self._maxOutputIndeces[batch, Y, X] = np.argmax(image[batch, y:y+self.maxPoolingSize, x:x+self.maxPoolingSize])
        return output
    
    def forwardOptim(self, image):
        """
        Es wird eine Kopie von der shape des Bild angefertigt, welche später im backward Pass verwendet wird.
        Wir "formen" den Image array um, sodass dieser aufgesplitet wird und entlang der Achsen -1, -2 den Imagearray
        in Form von kleinen "subArrays", die die einzelnen Teilausschnitte des Bildes in bezug auf die MaxPooling Operat-
        ion beinhalten.
        """
        self._input = image.shape
        _, y, x = image.shape
        newImage = image[:y - y % self.maxPoolingSize, :x - x % self.maxPoolingSize]
        newImage = newImage.reshape(-1, y // self.maxPoolingSize, self.maxPoolingSize, x // self.maxPoolingSize, self.maxPoolingSize)
        maxValuesImage = newImage.max(axis = (-3, -1))
        self._maxOutputIndeces1 = newImage.argmax(axis = -3)#this line makes me wanna give up
        return maxValuesImage
    
    def backward(self, gradient):
        dInput = np.zeros(self._input)
        b, h, w = self._input
        B, Y, X = np.mgrid[0:b, 0:h - h % self.maxPoolingSize:self.maxPoolingSize, 0:w - w % self.maxPoolingSize:self.maxPoolingSize]
        Y //= self.maxPoolingSize
        X //= self.maxPoolingSize
        maxIndex = np.unravel_index(self._maxOutputIndeces[B,Y, X].astype(int), (b,self.maxPoolingSize, self.maxPoolingSize))
        dInput[maxIndex[0], Y * self.maxPoolingSize + maxIndex[1], X * self.maxPoolingSize + maxIndex[2]] = gradient[B, Y, X]
        return dInput
    
    def backwardOptim(self, gradient):
        raise NotImplementedError

I tried pretty much everything i could think of. So yeah I don't know what to put here apart from my test setup to make it easier for you guys to help me.

if __name__ == "__main__":
    test = np.array([[[1, 2, 3, 4],
                      [5, 6, 7, 8],
                      
                      [1, 2, 3, 4],
                      [5, 6, 7, 8]],
                     
                     [[1, 2, 3, 4],
                      [5, 6, 7, 8],
                      [1, 2, 3, 4],
                      [5, 6, 7, 8]]])
    gradient = np.array([[[1, 2],
                         [1, 2]],
                         [[1, 2],
                          [1, 2]]])
    m = MaxPooling()
    output = m.forwardOptim(test)
    output1 = m.forward(test)
    dInput = m.backward(gradient)