Does propagate work only with 2 rank tensor inputs?

[Nokimann]

The code without using MessagePassing works.

I got an error message.

...
File "/home/noki/aifactory/modules.py", line 262, in message
  rm0 = F.silu(self.rbf_lin0(env*self.rbf0(r))).view(-1, self.num_heads, self.head_channels)
File "/home/noki/anaconda3/envs/gnn/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
  return forward_call(*input, **kwargs)
File "/home/noki/aifactory/modules.py", line 30, in forward
  denom = torch.where(r == 0., torch.tensor(1., device=r.device), r)

RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

edge_index: torch.Size([2, 36614])
query: torch.Size([2322, 8, 16])
key: torch.Size([2322, 8, 16])
value: torch.Size([2322, 8, 48])
_v: torch.Size([2322, 3, 8, 16])
r: torch.Size([36614])
_ur: torch.Size([36614, 3])
env: torch.Size([36614, 1])

        sm1, _vm3 = self.propagate(
            edge_index=edge_index, query=query, key=key, value=value, _v=_v, r=r, _ur=_ur, env=env
        )

    def message(self, query_i, key_j, value_j, _v_j, r, _ur, env):
        rm0 = F.silu(self.rbf_lin0(env*self.rbf0(r))).view(-1, self.num_heads, self.head_channels)
        rm1 = F.silu(self.rbf_lin1(env*self.rbf1(r))).view(-1, self.num_heads, 3*self.head_channels)
        attn = F.silu((query_i * key_j * rm0).sum(dim=2)) * env

        sm, vm0, vm1 = torch.split(value_j * rm1, self.head_channels, dim=2)
        
        sm = sm * attn.unsqueeze(2)
        _vm = _v_j * vm0.unsqueeze(1) + _ur.unsqueeze(2).unsqueeze(3) * vm1.unsqueeze(1)
        return sm, _vm

    def aggregate(self, inputs, index):
        sm, _vm = inputs
        sm1 = scatter_add(sm, index, dim=0)
        _vm3 = scatter_add(_vm, index, dim=0)
        return sm1, _vm3

    def update(self, inputs):
        return inputs

[rusty1s]

self.propagate works on tensors with any dimension. Importantly though, by default we expect that the node dimension is given in dimension -2, so you may want to fix this in your case:

class MyConv(MessagePassing):
     def __init__(self, ...):
          super().__init__(aggr=..., node_dim=0)

[Nokimann]

The code above uses the customized aggregate and update functions. Does that bother the aggr='add' in the super().__init__(...) or something? Or, how could I define the aggr argument and the customized aggregate function nicely?

[rusty1s]

No, it does not do any harm, but if you override aggregate, you can also pass aggr=None to the constructor.

[Nokimann]

Solved! Thank you always.