Segment_sum

[swamidass]

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

[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?

[fferflo]

There is currently no max_at and min_at in einx. However, you can use an einx adapter to implement this operation yourself. Here is an example with 1D indexing:

def max_at(tensor, indices, updates):
    return tensor.at[indices].max(updates)

ein_max_at = einx.jax.adapt_with_vmap(max_at)

tensor = jnp.zeros((10, 10))
indices = jnp.arange(5)
updates = jnp.ones((5,))

output = ein_max_at("a [b], [i], [i] -> a [b]", tensor, indices, updates)

And here is an example with 2D indexing:

def max_at(tensor, indices, updates):
    return tensor.at[indices[:, 0], indices[:, 1]].max(updates)

ein_max_at = einx.jax.adapt_with_vmap(max_at)

tensor = jnp.zeros((3, 10, 10))
coords = einx.id("a, a -> a (1 + 1)", jnp.arange(5), jnp.arange(5)) # Generate some 2D coordinates
updates = jnp.ones((5,))

output = ein_max_at("x [a b], [i 2], [i] -> [a b] x", tensor, coords, updates)

Closing this issue, feel free to reopen if anything isn't working as expected.