I went through the documentation of tf.nn.max_pool_with_argmax where it is written
Performs max pooling on the input and outputs both max values and indices.
The indices in argmax are flattened, so that a maximum value at position [b, y, x, c] becomes flattened index ((b * height + y) * width + x) * channels + c.
The indices returned are always in [0, height) x [0, width) before flattening, even if padding is involved and the mathematically correct answer is outside (either negative or too large). This is a bug, but fixing it is difficult to do in a safe backwards compatible way, especially due to flattening.
The variables b, y, x and c haven't been explicitly defined hence I was having issues implementing this method. Can someone please provide the same.