Absolute Max Aggregation Strategy

[baon6052]

Hi,

I would like to implement a custom aggregation strategy which takes the absolute maximum. I was wondering if this is already supported or in development? If not, would it be possible to give me some pointers on how I might implement this in the most efficient way?

I understand I would have to first perform scatter_max on the absolute values of my input tensor. Then I would need to map the argmax back to the original tensor, but in a way where I can preserve or retrieve the sign.

Thanks!

[rusty1s]

Why not do scatter_max(x.abs(), index)?

[rusty1s]

Oh, you can just do

out, argmax = scatter_max(x, index, ...)
out = x.scatter(argmax, dim=0)

[baon6052]

Thanks! Apologies, I may be misunderstanding how scatter works. Here is an example of my inputs:

Notice how the first element of the node at index 1 is 9.9000e+01,

inputs = tensor([
        [-1.9566e-02,  1.8216e-02, -1.2862e-02,  -5.8912e-03, 9.5555e-03,  3.9447e-02],
        [ 9.9000e+01,  1.8216e-02, -1.2862e-02, -5.8912e-03, 9.5555e-03,  3.9447e-02]]
        ...,
        grad_fn=<CopySlices>)

With a Star topology, where every node points only to node 0. After performing torch_scatter I get:

abs_max_values, argmax_indices = torch_scatter.scatter_max(inputs.abs(), index, dim=self.node_dim, dim_size=dim_size)

abs_max_values
tensor([
        [9.9000e+01, 1.8216e-02, 1.2862e-02,  5.8912e-03, 9.5555e-03, 3.9447e-02],
       [0.0000e+00, 0.0000e+00, 0.0000e+00,  0.0000e+00, 0.0000e+00, 0.0000e+00]
        ...,] grad_fn=<CppNode<ScatterMax>>)

argmax_indices
tensor([
        [ 1,  0,  0,  0,  0,  0],
        [11, 11, 11,  11, 11, 11],
         ...,])

Executing inputs.scatter(index=argmax_indices, dim=0, src=inputs) results in: *** RuntimeError: index 11 is out of bounds for dimension 0 with size 11.

Other useful info:

inputs.shape: torch.Size([11, 15452])
abs_max_values.shape: torch.Size([11, 15452])

Am I using argmax_indices in the correct way? Thanks!

[baon6052]

With inputs.scatter(index=argmax_indices, dim=-1, src=inputs), on a different set of inputs, this give me:

tensor([[ 1.7656e-03,  2.6976e-03,  3.3264e-03,  ...,  2.6048e-04,
         -5.6348e-03,  1.7656e-03],
        [ 9.9000e+01,  0.0000e+00,  0.0000e+00,  ...,  0.0000e+00,
         -0.0000e+00,  0.0000e+00]
        ...])

Which is closer to what I want. But the node at index 0 has not received the larger 99 value from the node at index 1.

[rusty1s]

Sorry, it's not scatter but gather:

import torch
from torch_scatter import scatter_max

x = torch.tensor([[-3, 2], [2, -3], [-2, 1], [1, -2]], dtype=torch.float)
index = torch.tensor([0, 0, 1, 1])

out, argmax = scatter_max(x.abs(), index, dim=0)
print(argmax)

out = torch.gather(x, 0, argmax)
print(out)

[baon6052]

Thanks this works!