Vectorized insert and searchsort without vmap

[Methylamphetamine]

Hi All,

Greetings!

Just posting to ask if there is any native function that supports "insert along axis" or "search sorted along axis", without vmapping the existing searchsorted and insert.

The user case is that I have a sorted score array scores of size BT, and the values values corresponding to the score of size BTF, where T is the number of tokens, and B is the batch size. I would like to find the places of S element per batch of new_scores and new_values (BS, BSF) into the score and values. The pseudo code would be

indices = jnp.search_along_axis(scores, new_scores, axis=1)
updated_scores = jnp.insert_along_axis(scores, new_scores, axis=1)
updated_values =  jnp.insert_along_axis(values, new_values, axis=1)

Currently searchsorted only supports 1D array search and insert, or other dynamic update functions, only support sliced inputs. I was wonder if it is possible to have some functions that do non-sliced insertion.

[jakevdp]

Hi - it's a good question! Actually I suspect the best way to write this would be something like

all_scores = jnp.concatenate(scores, new_scores, axis=1)
all_values = jnp.concatenate(values, new_values, axis=1)

indices = jnp.argsort(all_scores, axis=1)
batch_indices = jnp.arange(indices.shape[0])[:, None]

updated_scores = all_scores[batch_indices, indices]
updated_values = all_values[batch_indices, indices]

The reason for this is that XLA has no searchsorted primitive, so it can be somewhat slow in practice. And by doing static concatenation rather than dynamic insertion, all the operations will be JIT compatible.

Hope that helps!

[Methylamphetamine]

Hi Jake, thank you for the suggestion -- it is indeed a workaround without vmap! A following up question is that, does jax.numpy.sort uses quicksort by default, as numpy.sort? As far as I know quicksort does not work well for an almost sorted array, as in this case.

[jakevdp]

I'm actually not sure what sorting algorithm is used under the hood; it's not documented at https://openxla.org/xla/operation_semantics#sort, which means that it's probably determined by backend-specific code. You could probably figure it out for CPU by digging into the source code at https://github.com/openxla/xla/

Even if it uses quicksort, I suspect this approach would be far faster than any insertion-based approach implemented in terms of available XLA primitives.

[Methylamphetamine]

Hi Jake, I wrote a notebook to verify and it is indeed true! Thanks again for your help.