Max pooling by vector norm

[mariogeiger]

Hi,

I am looking for a way to do max pooling using indices in jax.

Here is what I exactly want to do:

window_dimension = 2
stride = 2
input = [
  [1.0, 0.0, 1.0],
  [2.0, 2.0, 0.0],
  [3.0, 0.0, 1.0],
  [0.0, 1.0, 1.0],
]
# In each window, pick the vector with the largest norm
output = [
  [2.0, 2.0, 0.0],
  [3.0, 0.0, 1.0],
]

[hawkinsp]

Perhaps something like this will work?

import jax.numpy as jnp
from jax import lax
import numpy as np


def normpool(x):
  norms = jnp.linalg.norm(x, axis=-1)
  idxs = jnp.arange(x.shape[0])

  def g(a, b):
    an, ai = a
    bn, bi = b
    which = an >= bn
    return (jnp.where(which, an, bn), jnp.where(which, ai, bi))

  _, idxs = lax.reduce_window((norms, idxs), (-np.inf, -1), g,
                    window_dimensions=(2,), window_strides=(2,),
                    padding=((0, 0),))
  return x[idxs]


input = jnp.array([
  [1.0, 0.0, 1.0],
  [2.0, 2.0, 0.0],
  [3.0, 0.0, 1.0],
  [0.0, 1.0, 1.0],
])
print(normpool(input))
# # In each window, pick the vector with the largest norm
# output = [
#   [2.0, 2.0, 0.0],
#   [3.0, 0.0, 1.0],
# ]

[mattjj]

It seems likely that the vmap rule wasn't updated to handle multiple array inputs!

[mariogeiger]

Shall I report a bug?

[hawkinsp]

@YouJiacheng I think variadic (here two-argument) reduce_window is only supported on CPU and TPU. It appears not to be implemented in XLA:GPU. One could work around inefficiently by extracting image patches and then performing a reduce instead of a reduce_window. I am confident that variadic non-windowed reduce is implemented everywhere.

For the vmap thing: yes, that's a JAX bug. It's probably easy to fix.

[YouJiacheng]

@hawkinsp But, why jit version works on GPU?

[mariogeiger]

I created an issue #9837