Scatter_logaddexp

[justinchiu]

Would it be possible to add support for the logaddexp binary operation?

[rusty1s]

So you mean a numerical stable version of scatter_add(src.exp(), index).log()? This should work just fine:

max_value = scatter_max(src, index, dim, dim_size)[0][index]
tmp = scatter_add((src - max_value).exp(), index, dim, dim_size)
out = torch.log(tmp + 1e-16) + max_value

If you know of a better way to implement it, please let me know :)

[justinchiu]

That should do it, thanks.
This is what I ended up with:

import torch
from torch_scatter import scatter_add, scatter_max

def scatter_lse(src, index, dim=-1, out=None, dim_size=None, fill_value=float("-inf")):
    dim_size = out.shape[dim] if dim_size is None and out is not None else dim_size
    max_value = scatter_max(src, index, dim=dim, dim_size=dim_size)[0]
    M = max_value.gather(dim, index)
    tmp = scatter_add(
        (src - M).exp(),
        index,
        dim=dim,
        out=out,
        dim_size=dim_size,
        fill_value=fill_value,
    )
    return torch.log(tmp + 1e-16) + max_value