Unpooling with FLAX

[clee1994]

Hey everyone,

I was wondering if there is a chance to realize an unpool operation? E.g. when I want to write a custom backward pass for the max pooling operation, how would I go about it? PyTorch has the option to return indices of the pooling operation and also provides an unpooling operation (https://pytorch.org/docs/stable/generated/torch.nn.MaxUnpool2d.html). Technically at some level, we must have access to the indices which are selected during max pool however I am unsure how I would access them.

Any help is appreciated!

Best,
Clemens