Repository: pyg-team/pytorch_geometric

Issue:
GATConv but for Dense Networks
## 🚀 Feature

Hi, I have a very simple request. I've seen that you've made dense formulations of a couple of the graph convolution types and graph pooling operators. Could you make the same thing for GATConv (and NNConv for that matter). If you don't have the bandwidth, just give me some pointers as to how to make this thing myself.
## Motivation

I'm basically using GATConv as a full self-attention network (all nodes connected).



Discussion:
That is a reasonable request, and I would be happy to include it into PyG. It would be great if you could work on this. For implementing, I would mostly follow the [`DenseGraphConv`](https://github.com/rusty1s/pytorch_geometric/blob/master/torch_geometric/nn/dense/dense_graph_conv.py) implementation. For GAT, you need to calculate attention weights first (with a masked softmax) before applying the matrix multiplication.
Hi there, my project also involves GAT operators on dense adjacency matrices. If this wasn't followed by anyone for now, I'll try to work on this for a `DenseGATConv` operator.
Super :)

Resolving patch:
diff --git a/torch_geometric/nn/dense/__init__.py b/torch_geometric/nn/dense/__init__.py
--- a/torch_geometric/nn/dense/__init__.py
+++ b/torch_geometric/nn/dense/__init__.py
@@ -1,4 +1,5 @@
 from .linear import Linear, HeteroLinear
+from .dense_gat_conv import DenseGATConv
 from .dense_sage_conv import DenseSAGEConv
 from .dense_gcn_conv import DenseGCNConv
 from .dense_graph_conv import DenseGraphConv
@@ -14,11 +15,12 @@
     'DenseGINConv',
     'DenseGraphConv',
     'DenseSAGEConv',
+    'DenseGATConv',
     'dense_diff_pool',
     'dense_mincut_pool',
     'DMoNPooling',
 ]
 
 lin_classes = __all__[:2]
-conv_classes = __all__[2:6]
-pool_classes = __all__[6:]
+conv_classes = __all__[2:7]
+pool_classes = __all__[7:]
diff --git a/torch_geometric/nn/dense/dense_gat_conv.py b/torch_geometric/nn/dense/dense_gat_conv.py
new file mode 100644
--- /dev/null
+++ b/torch_geometric/nn/dense/dense_gat_conv.py
@@ -0,0 +1,118 @@
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn import Parameter
+
+from torch_geometric.nn.dense.linear import Linear
+
+from ..inits import glorot, zeros
+
+
+class DenseGATConv(torch.nn.Module):
+    r"""See :class:`torch_geometric.nn.conv.GATConv`."""
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        heads: int = 1,
+        concat: bool = True,
+        negative_slope: float = 0.2,
+        dropout: float = 0.0,
+        bias: bool = True,
+    ):
+        # TODO Add support for edge features.
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.heads = heads
+        self.concat = concat
+        self.negative_slope = negative_slope
+        self.dropout = dropout
+
+        self.lin = Linear(in_channels, heads * out_channels, bias=False,
+                          weight_initializer='glorot')
+
+        # The learnable parameters to compute attention coefficients:
+        self.att_src = Parameter(torch.Tensor(1, 1, heads, out_channels))
+        self.att_dst = Parameter(torch.Tensor(1, 1, heads, out_channels))
+
+        if bias and concat:
+            self.bias = Parameter(torch.Tensor(heads * out_channels))
+        elif bias and not concat:
+            self.bias = Parameter(torch.Tensor(out_channels))
+        else:
+            self.register_parameter('bias', None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.lin.reset_parameters()
+        glorot(self.att_src)
+        glorot(self.att_dst)
+        zeros(self.bias)
+
+    def forward(self, x: Tensor, adj: Tensor, mask: Optional[Tensor] = None,
+                add_loop: bool = True):
+        r"""
+        Args:
+            x (torch.Tensor): Node feature tensor
+                :math:`\mathbf{X} \in \mathbb{R}^{B \times N \times F}`, with
+                batch-size :math:`B`, (maximum) number of nodes :math:`N` for
+                each graph, and feature dimension :math:`F`.
+            adj (torch.Tensor): Adjacency tensor
+                :math:`\mathbf{A} \in \mathbb{R}^{B \times N \times N}`.
+                The adjacency tensor is broadcastable in the batch dimension,
+                resulting in a shared adjacency matrix for the complete batch.
+            mask (torch.Tensor, optional): Mask matrix
+                :math:`\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}` indicating
+                the valid nodes for each graph. (default: :obj:`None`)
+            add_loop (bool, optional): If set to :obj:`False`, the layer will
+                not automatically add self-loops to the adjacency matrices.
+                (default: :obj:`True`)
+        """
+        x = x.unsqueeze(0) if x.dim() == 2 else x  # [B, N, F]
+        adj = adj.unsqueeze(0) if adj.dim() == 2 else adj  # [B, N, N]
+
+        H, C = self.heads, self.out_channels
+        B, N, _ = x.size()
+
+        if add_loop:
+            adj = adj.clone()
+            idx = torch.arange(N, dtype=torch.long, device=adj.device)
+            adj[:, idx, idx] = 1.0
+
+        x = self.lin(x).view(B, N, H, C)  # [B, N, H, C]
+
+        alpha_src = torch.sum(x * self.att_src, dim=-1)  # [B, N, H]
+        alpha_dst = torch.sum(x * self.att_dst, dim=-1)  # [B, N, H]
+
+        alpha = alpha_src.unsqueeze(1) + alpha_dst.unsqueeze(2)  # [B, N, N, H]
+
+        # Weighted and masked softmax:
+        alpha = F.leaky_relu(alpha, self.negative_slope)
+        alpha = alpha.masked_fill(adj.unsqueeze(-1) == 0, float('-inf'))
+        alpha = alpha.softmax(dim=2)
+        alpha = F.dropout(alpha, p=self.dropout, training=self.training)
+
+        out = torch.matmul(alpha.movedim(3, 1), x.movedim(2, 1))
+        out = out.movedim(1, 2)  # [B,N,H,C]
+
+        if self.concat:
+            out = out.reshape(B, N, H * C)
+        else:
+            out = out.mean(dim=2)
+
+        if self.bias is not None:
+            out = out + self.bias
+
+        if mask is not None:
+            out = out * mask.view(-1, N, 1).to(x.dtype)
+
+        return out
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.in_channels}, '
+                f'{self.out_channels}, heads={self.heads})')


Test patch:
diff --git a/test/nn/dense/test_dense_gat_conv.py b/test/nn/dense/test_dense_gat_conv.py
new file mode 100644
--- /dev/null
+++ b/test/nn/dense/test_dense_gat_conv.py
@@ -0,0 +1,66 @@
+import pytest
+import torch
+
+from torch_geometric.nn import DenseGATConv, GATConv
+from torch_geometric.testing import is_full_test
+
+
+@pytest.mark.parametrize('heads', [1, 4])
+@pytest.mark.parametrize('concat', [True, False])
+def test_dense_gat_conv(heads, concat):
+    channels = 16
+    sparse_conv = GATConv(channels, channels, heads=heads, concat=concat)
+    dense_conv = DenseGATConv(channels, channels, heads=heads, concat=concat)
+    assert str(dense_conv) == f'DenseGATConv(16, 16, heads={heads})'
+
+    # Ensure same weights and bias:
+    dense_conv.lin = sparse_conv.lin_src
+    dense_conv.att_src = sparse_conv.att_src
+    dense_conv.att_dst = sparse_conv.att_dst
+    dense_conv.bias = sparse_conv.bias
+
+    x = torch.randn((5, channels))
+    edge_index = torch.tensor([[0, 1, 1, 2, 3, 4], [1, 0, 2, 1, 4, 3]])
+
+    sparse_out = sparse_conv(x, edge_index)
+
+    x = torch.cat([x, x.new_zeros(1, channels)], dim=0).view(2, 3, channels)
+    adj = torch.Tensor([
+        [
+            [0, 1, 0],
+            [1, 0, 1],
+            [0, 1, 0],
+        ],
+        [
+            [0, 1, 0],
+            [1, 0, 0],
+            [0, 0, 0],
+        ],
+    ])
+    mask = torch.tensor([[1, 1, 1], [1, 1, 0]], dtype=torch.bool)
+
+    dense_out = dense_conv(x, adj, mask)
+
+    assert dense_out[1, 2].abs().sum() == 0
+    dense_out = dense_out.view(6, dense_out.size(-1))[:-1]
+    assert torch.allclose(sparse_out, dense_out, atol=1e-4)
+
+    if is_full_test():
+        jit = torch.jit.script(dense_conv)
+        assert torch.allclose(jit(x, adj, mask), dense_out)
+
+
+def test_dense_gat_conv_with_broadcasting():
+    batch_size, num_nodes, channels = 8, 3, 16
+    conv = DenseGATConv(channels, channels, heads=4)
+
+    x = torch.randn(batch_size, num_nodes, channels)
+    adj = torch.Tensor([
+        [0, 1, 1],
+        [1, 0, 1],
+        [1, 1, 0],
+    ])
+
+    assert conv(x, adj).size() == (batch_size, num_nodes, 64)
+    mask = torch.tensor([1, 1, 1], dtype=torch.bool)
+    assert conv(x, adj, mask).size() == (batch_size, num_nodes, 64)
diff --git a/test/nn/dense/test_dense_gcn_conv.py b/test/nn/dense/test_dense_gcn_conv.py
--- a/test/nn/dense/test_dense_gcn_conv.py
+++ b/test/nn/dense/test_dense_gcn_conv.py
@@ -10,7 +10,7 @@ def test_dense_gcn_conv():
     dense_conv = DenseGCNConv(channels, channels)
     assert str(dense_conv) == 'DenseGCNConv(16, 16)'
 
-    # Ensure same weights and bias.
+    # Ensure same weights and bias:
     dense_conv.lin.weight = sparse_conv.lin.weight
     dense_conv.bias = sparse_conv.bias
 
@@ -39,14 +39,14 @@ def test_dense_gcn_conv():
     dense_out = dense_conv(x, adj, mask)
     assert dense_out.size() == (2, 3, channels)
 
+    assert dense_out[1, 2].abs().sum() == 0
+    dense_out = dense_out.view(6, channels)[:-1]
+    assert torch.allclose(sparse_out, dense_out, atol=1e-4)
+
     if is_full_test():
         jit = torch.jit.script(dense_conv)
         assert torch.allclose(jit(x, adj, mask), dense_out)
 
-    assert dense_out[1, 2].abs().sum().item() == 0
-    dense_out = dense_out.view(6, channels)[:-1]
-    assert torch.allclose(sparse_out, dense_out, atol=1e-04)
-
 
 def test_dense_gcn_conv_with_broadcasting():
     batch_size, num_nodes, channels = 8, 3, 16