How should I do segment_argmax in JAX?

[ameya98]

I need to identify the indices corresponding to the largest values in segments of an array. Basically, the argmax version of jax.ops.segment_max. How would one implement this?

[jakevdp]

Very good question – I don't think there's a good answer at the moment. jax.ops.segment_* is not implemented with variadic inputs in mind, which is what we'd need in order to do segment_argmax.

[jakevdp]

Here's a basic implementation you might use, though if you have many segments it's not particularly efficient:

import jax
import numpy as np
import jax.numpy as jnp

def segment_argmax(data, segment_ids, num_segments=None):
  if num_segments is None:
    num_segments = np.max(segment_ids) + 1
  num_segments = int(num_segments)
  data = jnp.asarray(data)
  segment_ids = jnp.asarray(segment_ids)
  return jax.vmap(lambda i: jnp.where(i == segment_ids, data, -jnp.inf))(
      jnp.arange(num_segments)).argmax(1)

data = jnp.array(np.random.randint(0, 100, 10))
segments = jnp.array([0, 1, 2, 0, 2, 2, 1, 2, 0, 1])

print("data:", data)
print("segments:", segments)
print(jax.ops.segment_max(data, segments))
print(segment_argmax(data, segments))

data: [16 37 23 68 97 69 85 10 15 96]
segments: [0 1 2 0 2 2 1 2 0 1]
[68 96 97]
[3 9 4]

[ameya98]

Thank you so much! This works perfectly.