Memory-Efficient `logsumexp` Over Unequal Partitions in JAX #28046

Qazalbash · 2025-04-15T21:17:58Z

Qazalbash
Apr 15, 2025

Hi,

I’m working on a JAX project and encountered a challenge related to memory efficiency. I have a large 1D array of shape (N,), which is partitioned into d contiguous segments of known, but unequal, sizes: (n₁,), (n₂,), ..., (n_d,), such that they form a complete partition of the original array.

My goal is to compute the logsumexp over each subarray efficiently. While the shapes of the segments are known at compile time, the unequal sizes seem to prevent fully vectorized operations, and the current implementation consumes significant memory.

Could you advise on the most memory-efficient approach to compute the logsumexp over these partitions in JAX? I'm particularly interested in strategies that avoid materializing padded arrays or large intermediate representations.

Any insights or suggestions would be greatly appreciated.

import jax
from jax import Array, nn as jnn, numpy as jnp, random as jrd


# Compute start/stop indices for each event from data_shapes
# fmt: off
data_shapes = [
    10856, 4379, 10224, 12001, 12001, 12001, 11735, 12001, 11844, 10596, 11396, 11939,
    11982, 11724, 12001, 5223, 6438, 7503, 12001, 728, 12001, 8218, 4084, 11982, 8171,
    12001, 12001, 12001, 5596, 3603, 7721, 12001, 11995, 12001, 3835, 1191, 3465, 12001,
    12001, 12001, 12001, 12001, 12001, 3044, 12001, 10890, 5422, 11956, 12001, 12001,
    11954, 2190, 3882, 12001, 11433, 12001, 2235, 11976, 3384, 10755, 12001, 8652, 6599,
    12001, 12001, 10918, 2577, 12001, 11685, 11967, 11047, 11528, 12001, 12001, 12001,
    4148, 11539, 4107, 12001, 12001, 7697, 6040, 12001, 1743, 11995, 12001, 12001, 4885,
    10766, 2508, 11902, 7059, 12001, 12001, 7799, 6747, 10991, 12001, 1600, 12001,
    12001, 12001, 12001, 11892, 12001, 12001, 2908, 12001, 1917, 10731, 10971, 12001,
    5518, 12001, 11965, 11108, 12001, 8436, 11959, 5483, 11992, 5224, 1151
]
# fmt: on
start_indices = [0] + list(jnp.cumsum(jnp.array(data_shapes[:-1])).tolist())
stop_indices = [start + size for start, size in zip(start_indices, data_shapes)]
start_stops = list(zip(start_indices, stop_indices))

big_array = jrd.normal(key=jrd.PRNGKey(0), shape=(sum(data_shapes), 5))


def f(x: Array) -> Array:
    # Simulate some computation
    # In practice, its an expensive function
    return jnp.sum(x, axis=-1)


def g(x: Array) -> Array:
    return jnp.sum(x, axis=-1)


@jax.jit
def main_fn(x: Array) -> Array:
    g_x = g(x)
    f_big_arrray = jax.block_until_ready(jax.lax.map(f, big_array, batch_size=10_000))

    total = jnp.zeros(())
    for start, stop in start_stops:
        f_big_arrray_slice = jax.lax.dynamic_slice_in_dim(
            f_big_arrray, start, stop - start, axis=-1
        )
        total += jnn.logsumexp(f_big_arrray_slice, axis=-1) - jnp.log(stop - start)

    return g_x + total


if __name__ == "__main__":
    # Run the main function
    xx = jrd.normal(key=jrd.PRNGKey(37), shape=(10, 3))
    # vmap over the main function
    # it is in other lib which i can not change
    # I can only change the implementation of main_fn
    result = jax.vmap(main_fn)(xx)

Thank you

jakevdp · 2025-04-15T21:25:28Z

jakevdp
Apr 15, 2025
Maintainer

It sounds like what you want is a segment_logsumexp, with semantics similar to the existing segment_sum operation.

You could probably write this function yourself by scaling the full array by the max value, then computing the exp, the segment_sum the log, and then rescaling the result: effectively you'd be re-implementing logsumexp, but replacing sum with segment_sum. What do you think?

7 replies

Qazalbash Apr 15, 2025
Author

@jakevdp This implementation is exhausting the resources. Any tips to optimize it? Before that is this correct or not?

jakevdp Apr 15, 2025
Maintainer

The implementation looks reasonable to me – I'd probably write a few tests to make sure it's behaving as expected.

Regarding exhausting of resources... how large is your input array?

Qazalbash Apr 16, 2025
Author

@jakevdp

The implementation looks reasonable to me – I'd probably write a few tests to make sure it's behaving as expected.

I will write them.

Regarding exhausting of resources... how large is your input array?

They are in orders of 10^6 or greater. As I mentioned the function which will use segment_logsumexp will be passed into a vmap. I think this is increasing the demand for memory.

Qazalbash Apr 17, 2025
Author

My understanding is that this function should become fused,

Code

from functools import partial

import jax
from jax import Array, lax, numpy as jnp
from jax._src.numpy.util import promote_args_inexact
from jaxtyping import ArrayLike


# fmt: off
sizes = [
    10856, 4379, 10224, 12001, 12001, 12001, 11735, 12001, 11844, 10596, 11396, 11939,
    11982, 11724, 12001, 5223, 6438, 7503, 12001, 728, 12001, 8218, 4084, 11982, 8171,
    12001, 12001, 12001, 5596, 3603, 7721, 12001, 11995, 12001, 3835, 1191, 3465, 12001,
    12001, 12001, 12001, 12001, 12001, 3044, 12001, 10890, 5422, 11956, 12001, 12001,
    11954, 2190, 3882, 12001, 11433, 12001, 2235, 11976, 3384, 10755, 12001, 8652, 6599,
    12001, 12001, 10918, 2577, 12001, 11685, 11967, 11047, 11528, 12001, 12001, 12001,
    4148, 11539, 4107, 12001, 12001, 7697, 6040, 12001, 1743, 11995, 12001, 12001, 4885,
    10766, 2508, 11902, 7059, 12001, 12001, 7799, 6747, 10991, 12001, 1600, 12001,
    12001, 12001, 12001, 11892, 12001, 12001, 2908, 12001, 1917, 10731, 10971, 12001,
    5518, 12001, 11965, 11108, 12001, 8436, 11959, 5483, 11992, 5224, 1151
]
# fmt: on


@partial(
    jax.jit,
    static_argnames=(
        "num_segments",
        "indices_are_sorted",
        "unique_indices",
        "bucket_size",
        "mode",
    ),
)
def segment_logsumexp(
    data: ArrayLike,
    segment_ids: ArrayLike,
    num_segments: int | None = None,
    indices_are_sorted: bool = False,
    unique_indices: bool = False,
    bucket_size: int | None = None,
    mode: lax.GatherScatterMode | None = None,
) -> Array:
    (data,) = promote_args_inexact("segment_logsumexp", data)

    amax = jnp.max(data.real, initial=-jnp.inf)
    amax = lax.stop_gradient(
        lax.select(jnp.isfinite(amax), amax, lax.full_like(amax, 0))
    )

    exp_a = lax.exp(lax.sub(data, amax))

    sumexp = jax.ops.segment_sum(
        exp_a,
        segment_ids,
        num_segments=num_segments,
        indices_are_sorted=indices_are_sorted,
        unique_indices=unique_indices,
        bucket_size=bucket_size,
        mode=mode,
    )

    out = lax.add(lax.log(sumexp), amax.astype(sumexp.dtype))

    return out


def f(x: Array) -> Array:
    return jnp.square(x)


data = jnp.zeros((sum(sizes),))
segement_ids = jnp.concatenate([jnp.full((size,), i) for i, size in enumerate(sizes)])


print(
    jax.jit(
        segment_logsumexp,
        static_argnames=(
            "num_segments",
            "indices_are_sorted",
            "unique_indices",
            "bucket_size",
            "mode",
        ),
    )
    .lower(data, segement_ids, len(sizes), False, False, None, None)
    .compile()
    .as_text()
)

HLO Text

HloModule jit_segment_logsumexp, is_scheduled=true, entry_computation_layout={(f32[1159351]{0}, s32[1159351]{0})->f32[123]{0}}, allow_spmd_sharding_propagation_to_parameters={true,true}, allow_spmd_sharding_propagation_to_output={true}, frontend_attributes={fingerprint_before_lhs="de21e4be8a8ba7ed196142e9da83c486"}

%region_0.3 (Arg_0.4: f32[], Arg_1.5: f32[]) -> f32[] {
  %Arg_1.5 = f32[] parameter(1), metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/reduce_max"}
  %Arg_0.4 = f32[] parameter(0), metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/reduce_max"}
  ROOT %maximum.3 = f32[] maximum(%Arg_0.4, %Arg_1.5), metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/reduce_max" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=122}
}

%fused_reduce (param_0.9: f32[1159351]) -> f32[1252] {
  %param_0.9 = f32[1159351]{0} parameter(0)
  %constant_1_1 = f32[] constant(-inf)
  %pad.1.1 = f32[1159352]{0} pad(%param_0.9, %constant_1_1), padding=0_1, metadata={op_name="data"}
  %bitcast.106.1 = f32[1252,926]{1,0} bitcast(%pad.1.1)
  ROOT %reduce.1.1 = f32[1252]{0} reduce(%bitcast.106.1, %constant_1_1), dimensions={1}, to_apply=%region_0.3, metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/reduce_max" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=122}
}

%fused_reduce.1 (param_0.10: f32[1252]) -> f32[] {
  %param_0.10 = f32[1252]{0} parameter(0)
  %constant_1_2 = f32[] constant(-inf)
  ROOT %reduce.0.1 = f32[] reduce(%param_0.10, %constant_1_2), dimensions={0}, to_apply=%region_0.3, metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/reduce_max" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=122}
}

%fused_slice (param_1_0: f32[]) -> (f32[123], f32[1]) {
  %constant_0_3 = f32[] constant(0)
  %broadcast.9.2 = f32[123]{0} broadcast(%constant_0_3), dimensions={}
  %param_1_0 = f32[] parameter(0)
  %is-finite.3.2 = pred[] is-finite(%param_1_0), metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/is_finite" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=124}
  %select.3.2 = f32[] select(%is-finite.3.2, %param_1_0, %constant_0_3), metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/select_n" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=124}
  %reshape.2 = f32[1]{0} reshape(%select.3.2), metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/select_n" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=124}
  %concatenate = f32[124]{0} concatenate(%broadcast.9.2, %reshape.2), dimensions={0}
  %slice = f32[123]{0} slice(%concatenate), slice={[0:123]}
  %slice.1 = f32[1]{0} slice(%concatenate), slice={[123:124]}
  ROOT %tuple = (f32[123]{0}, f32[1]{0}) tuple(%slice, %slice.1), metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/select_n" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=124}
}

%region_1.7 (Arg_0.8: f32[], Arg_1.9: f32[]) -> f32[] {
  %Arg_1.9 = f32[] parameter(1), metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/scatter-add"}
  %Arg_0.8 = f32[] parameter(0), metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/scatter-add"}
  ROOT %add.5 = f32[] add(%Arg_0.8, %Arg_1.9), metadata={op_name="/add" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=129}
}

%fused_scatter (param_0: f32[123], param_1.4: s32[1159351], param_2.4: f32[1159351], param_3.1: f32[]) -> f32[123] {
  %param_0 = f32[123]{0} parameter(0)
  %param_1.4 = s32[1159351]{0} parameter(1)
  %bitcast.97.1 = s32[1159351,1]{1,0} bitcast(%param_1.4)
  %param_2.4 = f32[1159351]{0} parameter(2)
  %param_3.1 = f32[] parameter(3)
  %broadcast.10.1 = f32[1159351]{0} broadcast(%param_3.1), dimensions={}, metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/sub" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=127}
  %subtract.3.1 = f32[1159351]{0} subtract(%param_2.4, %broadcast.10.1), metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/sub" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=127}
  %exponential.3.1 = f32[1159351]{0} exponential(%subtract.3.1), metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/exp" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=127}
  %bitcast.118.1 = f32[1159351,1]{1,0} bitcast(%exponential.3.1)
  ROOT %scatter.4.1 = f32[123]{0} scatter(%param_0, %bitcast.97.1, %bitcast.118.1), update_window_dims={1}, inserted_window_dims={}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_1.7, metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/scatter-add" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=129}
}

%fused_add (param_0.5: f32[123], param_1.8: f32[]) -> f32[123] {
  %param_0.5 = f32[123]{0} parameter(0)
  %log.3.1 = f32[123]{0} log(%param_0.5), metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/log" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=139}
  %param_1.8 = f32[] parameter(1)
  %broadcast.11.1 = f32[123]{0} broadcast(%param_1.8), dimensions={}, metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/add" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=139}
  ROOT %add.6.1 = f32[123]{0} add(%log.3.1, %broadcast.11.1), metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/add" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=139}
}

%command_buffer (p: f32[1159351], p.1: s32[1159351]) -> f32[123] {
  %p = f32[1159351]{0} parameter(0)
  %p.1 = s32[1159351]{0} parameter(1)
  %input_reduce_fusion = f32[1252]{0} fusion(%p), kind=kInput, calls=%fused_reduce, metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/reduce_max" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=122}
  %input_reduce_fusion.1 = f32[] fusion(%input_reduce_fusion), kind=kInput, calls=%fused_reduce.1, metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/reduce_max" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=122}
  %input_slice_fusion = (f32[123]{0}, f32[1]{0}) fusion(%input_reduce_fusion.1), kind=kInput, calls=%fused_slice, metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/select_n" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=124}
  %get-tuple-element.1 = f32[1]{0} get-tuple-element(%input_slice_fusion), index=1
  %bitcast.130 = f32[] bitcast(%get-tuple-element.1), metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/select_n" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=124}
  %get-tuple-element = f32[123]{0} get-tuple-element(%input_slice_fusion), index=0
  %input_scatter_fusion = f32[123]{0} fusion(%get-tuple-element, %p.1, %p, %bitcast.130), kind=kInput, calls=%fused_scatter, metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/scatter-add" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=129}
  ROOT %loop_add_fusion = f32[123]{0} fusion(%input_scatter_fusion, %bitcast.130), kind=kLoop, calls=%fused_add, metadata={op_name="jit(segment_logsumexp)/jit(main)/jit(segment_logsumexp)/add" source_file="/home/gradf/academia/gwkokab/playground.py" source_line=139}
}

ENTRY %main.31 (Arg_0.1.0: f32[1159351], Arg_1.2.0: s32[1159351]) -> f32[123] {
  %Arg_0.1.0 = f32[1159351]{0} parameter(0), metadata={op_name="data"}
  %Arg_1.2.0 = s32[1159351]{0} parameter(1), metadata={op_name="segment_ids"}
  ROOT %call = f32[123]{0} call(%Arg_0.1.0, %Arg_1.2.0), to_apply=%command_buffer
}

Qazalbash Apr 19, 2025
Author

I didn't find any way to reduce the memory consumption using the segment_logsumexp.

I padded the arrays and used the jax.lax.scan to work around. In my case, it is working better than ever. You can see the file gwkokab/inference/poissonlikelihood.py, reference PR gwkokab/gwkokab#526.

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Memory-Efficient `logsumexp` Over Unequal Partitions in JAX #28046

{{title}}

Replies: 1 comment 7 replies

{{title}}

{{title}}

{{title}}

{{title}}

{{editor}}'s edit

{{editor}}'s edit

{{title}}

{{title}}

{{editor}}'s edit

{{editor}}'s edit

Select a reply

Memory-Efficient logsumexp Over Unequal Partitions in JAX #28046

Qazalbash Apr 15, 2025

Replies: 1 comment · 7 replies

jakevdp Apr 15, 2025 Maintainer

Qazalbash Apr 15, 2025 Author

jakevdp Apr 15, 2025 Maintainer

Qazalbash Apr 16, 2025 Author

Qazalbash Apr 17, 2025 Author

Qazalbash Apr 19, 2025 Author

Memory-Efficient `logsumexp` Over Unequal Partitions in JAX #28046

Qazalbash
Apr 15, 2025

Replies: 1 comment 7 replies

jakevdp
Apr 15, 2025
Maintainer

Qazalbash Apr 15, 2025
Author

jakevdp Apr 15, 2025
Maintainer

Qazalbash Apr 16, 2025
Author

Qazalbash Apr 17, 2025
Author

Qazalbash Apr 19, 2025
Author