Repository: pyg-team/pytorch_geometric (python)

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).


Pull request description:
Add `DenseGATConv` operator
Added `DenseGATConv` operator that supports most features of the sparse counterpart except bipartite graphs (which do not makes sense given adjacency matrices) and edge attention.

This PR resolves #1141.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 67c67b9a6..c8f25bf2e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added the `DenseGATConv` module ([#6928](https://github.com/pyg-team/pytorch_geometric/pull/6928))
 - Added `trim_to_layer` utility function for more efficient `NeighborLoader` use-cases ([#6661](https://github.com/pyg-team/pytorch_geometric/pull/6661))
 - Added the `DistMult` KGE model ([#6958](https://github.com/pyg-team/pytorch_geometric/pull/6958))
 - Added `HeteroData.set_value_dict` functionality ([#6961](https://github.com/pyg-team/pytorch_geometric/pull/6961), [#6974](https://github.com/pyg-team/pytorch_geometric/pull/6974))
diff --git a/torch_geometric/nn/dense/__init__.py b/torch_geometric/nn/dense/__init__.py
index 3848b2038..c7ea2a982 100644
--- 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 @@ __all__ = [
     '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
index 000000000..20f90469a
--- /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
index 000000000..72c4ccc4b
--- /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
index 730f73e79..7d43ec0fc 100644
--- 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