How to return the indices of max values with max pooling?

[LauriKurki]

Hi!

Is there a way to return the indices of the maximum values returned with flax.linen.max_pool? Alternatively, how can I make use of the flax.linen.pool helper function in making my own pooling function?

For example,

import jax.numpy as jnp

x = jnp.array([
    [1, 2, 3, 4],
    [5, 6, 7, 8],
    [9, 10, 11, 12],
    [13, 14, 15, 16]
])
x = x[None, ..., None]

pooled_x, idxs = max_pool_with_indices(x, (2, 2), (2, 2), 'SAME')
print(idxs)

would output

[[[[ 5]
   [ 7]]

  [[13]
   [15]]]]

as these indices correspond to the maximum values in each window.

If it's any help, this is implemeted in Pytorch (https://pytorch.org/docs/stable/generated/torch.nn.MaxPool3d.html) with argument return_indices=True.

Best,
Lauri

[cherrywoods]

I don't know about flax, but you can use jax.lax.reduce_window to achieve this, as discussed here: jax-ml/jax#9818 (comment)