JIT compatible equivalent to segment_sum that concatenates instead of summing?

[NiklasKappel]

Hi, I am new to (thinking in) JAX and trying to implement a function segment_concatenate in analogy to jax.ops.segment_sum with the signature

def segment_concatenate(data, segment_ids, num_segments, max_duplicate_indices, padding_value):

where the output is an array x of shape (num_segments, max_duplicate_indices), and x[i] is the concatenation (instead of summation) of all values in data that correspond to the index i. If i appears in segment_ids less often than max_duplicate_indices times, the remaining spots should be filled by padding_value.

Is it possible to implement this in a JIT compatible way? The output shape is static, but I have not yet come up with a solution that avoids intermediate variable-shaped arrays.

[NiklasKappel]

This function does what I want to and is JIT compatible, but I am not sure if the combination of for loop and repeated .at().set() is efficient or idiomatic.

from functools import partial
from typing import Any

import jax
import jax.numpy as jnp
from jaxtyping import Array, Integer, Shaped


@partial(jax.jit, static_argnums=(2, 3))
def segment_concatenate(
    data: Shaped[Array, "E"],
    segment_ids: Integer[Array, "E"],
    num_segments: int,
    max_duplicate_indexes: int,
    padding_value: Any,
) -> Shaped[Array, "{num_segments} {max_duplicate_indexes}"]:
    output = jnp.full(
        (num_segments, max_duplicate_indexes), padding_value, dtype=data.dtype
    )
    segment_id_counts = jnp.zeros(num_segments, dtype=jnp.int32)
    for i, segment_id in enumerate(segment_ids):
        output = output.at[segment_id, segment_id_counts[segment_id]].set(data[i])
        segment_id_counts = segment_id_counts.at[segment_id].add(1)
    return output

[jakevdp]

This looks like a decent enough approach, and should be efficient so long as segment_ids is quite small. If segment_ids is longer, the flattening of the python for loop could lead to prohibitively long compilation times – I'd suggest rewriting that part in a vectorized manner if possible. It seems like it would effectively be a generalization of some of the approaches used in the implementation of jnp.unique.

[NiklasKappel]

I settled on this for now, which has a scan instead of the for loop.

@partial(jax.jit, static_argnums=(2, 3))
def segment_concatenate(
    data: Shaped[Array, " E"],
    segment_ids: Integer[Array, " E"],
    num_segments: int,
    max_duplicate_indexes: int,
    padding_value: Any,
) -> Shaped[Array, "{num_segments} {max_duplicate_indexes}"]:
    output = jnp.full(
        (num_segments, max_duplicate_indexes), padding_value, dtype=data.dtype
    )

    if segment_ids.shape[0] == 0:
        return output

    segment_id_counts = jnp.zeros(num_segments, dtype=jnp.int32)

    def update_fn(carry, idx):
        counts, out = carry
        segment_id = segment_ids[idx]
        out = out.at[segment_id, counts[segment_id]].set(data[idx])
        counts = counts.at[segment_id].add(1)
        return (counts, out), None

    (_, output), _ = jax.lax.scan(
        update_fn, (segment_id_counts, output), jnp.arange(data.shape[0])
    )

    output = jnp.sort(output)
    return output

In contrast to segment_sum, it considers negative segment ids as valid (wrapping -1 to the last element etc.).

[jakevdp]

Here's a vectorized approach that does more-or-less the same thing, except it pads at the end instead of at the beginning, and it deals differently with negative indices:

def segment_concatenate(
    data: jax.Array,
    segment_ids: jax.Array,
    num_segments: int,
    max_duplicate_indexes: int,
    padding_value: Any,
) -> jax.Array:
  vals, indices = jnp.unique(segment_ids, return_inverse=True, size=num_segments, fill_value=-1)
  mask = vals[:, None] == indices[None, :]
  i, j = jnp.where(mask, size=num_segments * max_duplicate_indexes, fill_value=max_duplicate_indexes)
  k = jnp.cumsum(mask, axis=1).at[i, j].get(fill_value=max_duplicate_indexes) - 1
  return jnp.zeros_like(data, shape=(num_segments, max_duplicate_indexes)).at[i, k].set(data[j])

data = jnp.arange(10, 20)
segment_ids = jnp.array([0, 1, 2, 0, 2, 1, 1, 2, 2, 2])
segment_concatenate(data, segment_ids, num_segments=3, max_duplicate_indexes=5, padding_value=0)
# Array([[10, 13,  0,  0,  0],
#        [11, 15, 16,  0,  0],
#        [12, 14, 17, 18, 19]], dtype=int32)

Not comprehensively tested, but hopefully it's useful as an idea for an alternate approach.

[NiklasKappel]

Thank you very much, it was very instructive to work through this! I made some small changes (/corrections):

def segment_concatenate(
    data: jax.Array,
    segment_ids: jax.Array,
    num_segments: int,
    max_duplicate_indexes: int,
    padding_value: Any,
) -> jax.Array:
    vals, indices = jnp.unique(
        segment_ids, return_inverse=True, size=num_segments, fill_value=-1
    )
    mask = vals[:, None] == vals[indices][None, :]
    i, j = jnp.where(mask, size=len(segment_ids))
    k = jnp.cumsum(mask, axis=1).at[i, j].get() - 1
    return (
        jnp.full_like(data, padding_value, shape=(num_segments, max_duplicate_indexes))
        .at[i, k]
        .set(data[j])
    )


data = jnp.arange(10, 20)
segment_ids = jnp.array([3, 1, 2, 3, 2, 1, 1, 2, 2, 2])
segment_concatenate(
    data, segment_ids, num_segments=3, max_duplicate_indexes=5, padding_value=-1
)
# Array([[11, 15, 16, -1, -1],
#        [12, 14, 17, 18, 19],
#        [10, 13, -1, -1, -1]], dtype=int32)

• I use vals[indices][None, :] to handle segment_ids that are not contiguous starting from 0 (I guess that's why we use jnp.unique in the first place).
• I think len(segment_ids) is the correct size for i,j, because that the number of True entries in the mask (each data entry goes to exactly one segment).
• With this we don't need fill values in mask and get.
• At the end, we use full_like with the padding_value.