Segment_sum

Open swamidass opened this issue 7 months ago • 1 comments

Is there a way to do the equivalent of segment_sum (and product, softmax, etc) in einx?

https://docs.jax.dev/en/latest/_autosummary/jax.ops.segment_sum.html jax.ops.segment_sum — JAX documentation

Jun 05 '25 00:06 swamidass

I found that there is a way to do this with einx using "add_at".

But to correctly implement logmax, we also need the equivalent of "max_at" or "min_at". Can this be added in?

Jun 18 '25 16:06 swamidass