How to efficiently implement large segment sums

[benjaminvatterj]

Hi!

I have a model that requires summing some output over a very large collection of groups, like what jax.ops.segement_sum does. The segment indices are model constants, which trigger a very slow and painful constant folding when compiling, which I would like to avoid. However, I have yet to find an alternative to segment_sum that isn't very slow or results in a segmentation fault.

Here is a minimal reproducible example of the problem. There are three version of the problem, the first is the original that works really well except when the derivative is called. The second is too slow to use in practice, and the third results in segmentation fault.

import jax.numpy as jnp
import jax

model_data = {
    'num_segments': 11544,
    'segment_id': jnp.repeat(jnp.arange(11544), 30),
    'val_id': jnp.tile(jnp.arange(10), 3 * 11544),
}


# This one works very fast (1.28 ms), but results in a slow constant folding especially when calling the jvp
@jax.jit
def original(x):
    segment_id = model_data['segment_id']
    val_id = model_data['val_id']
    num_segment = model_data['num_segments']
    return jax.ops.segment_sum(x[val_id], segment_id, num_segment)

# This one works slow (1.41 s), but seems to avoid the constant folding issue on the jvp
@jax.jit
def slow_version(x):
    segment_id = model_data['segment_id']
    val_id = model_data['val_id']
    num_segments = model_data['num_segments']
    agg_mat = (jnp.arange(num_segments).reshape(-1, 1) == segment_id.reshape(1, -1))
    x = x[val_id]
    return agg_mat @ x

# This one results in segfault
@jax.jit
def segfault_version(x):
    with jax.ensure_compile_time_eval():
        segment_id = model_data['segment_id']
        val_id = model_data['val_id']
        num_segments = model_data['num_segments']
        agg_mat = (jnp.arange(num_segments).reshape(-1, 1) == segment_id.reshape(1, -1))
    x = x[val_id]
    return agg_mat @ x

x = jnp.ones(5)
print(original(x))

jax version: 0.4.20 on CPU (mac studio M2 ultra)
any help would be greatly appreciated!

[Qazalbash]

@benjaminvatterj Hi, did you find any solution?

[benjaminvatterj]

No. I ended up specifying the jvp to reduce the cost slightly, but it is still very demanding. In general, I find that working with Jax on large data models results in significant overhead due to constant folding, which is offset by much faster execution down the line.