Thanks a lot for following up!
Yes, my kernel not only varies in size, but also operates on non-adjacent neurons.
In essence, I’m trying to do the following: given a layer taking as input a tensor of the shape (3, 10, 10), I flatten it (i.e. input shape is a tensor of shape (75)), then throw away some neurons because I deem them unnecessary, and then I want to perform the “partial” max pooling on the result. I can use unfold to figure out how I would have slid the kernel on the flattened layer, however, if some of the neurons are no longer present in the layer, I do not know how to do the MaxPooling on the remaining.
Maybe my earlier example was a bit too hand-wavy. Let me try to be a little more concrete.
Suppose I have the following input
[[a0, a1, a2]
 [a3, a4, a5]
 [a6, a7, a8]]

MaxPooling with a kernel of size (2,2) will produce the max over the following windows
[[a0, a1]
 [a3, a4]]

[[a1, a2]
 [a4, a5]]

[[a3, a4]
 [a6, a7]]

[[a4, a5]
 [a7, a8]]

Now suppose, I had flattened my input
[a0, a1, a2, a3, a4, a5, a6, a7, a8]

Now I can think of the windows as following
[a0, a1, a3, a4]

[a1, a2, a4, a5]

[a3, a4, a6, a7]

[a4, a5, a7, a8]

Suppose, I deemed that neuron which processes a3 was unnecessary. So I prune it away while keeping all the other layers intact. Now you can imagine how the MaxPooling would change. The windows would be
[a0, a1, a4] instead of [a0, a1, a3, a4]
[a1, a2, a4, a5] remains the same.
[a4, a6, a7] instead of [a3, a4, a6, a7]
[a4, a5, a7, a8] remains the same.