update

test/utils/test_to_dense_adj.py

@@ -23,6 +23,11 @@ def test_to_dense_adj():
         assert adj[0].tolist() == [[1, 1, 0], [1, 0, 0], [0, 0, 0]]
         assert adj[1].tolist() == [[0, 1, 0], [0, 0, 1], [1, 0, 0]]
 
+    adj = to_dense_adj(edge_index, batch, max_num_nodes=2)
+    assert adj.size() == (2, 2, 2)
+    assert adj[0].tolist() == [[1, 1], [1, 0]]
+    assert adj[1].tolist() == [[0, 1], [0, 0]]
+
     adj = to_dense_adj(edge_index, batch, max_num_nodes=5)
     assert adj.size() == (2, 5, 5)
     assert adj[0][:3, :3].tolist() == [[1, 1, 0], [1, 0, 0], [0, 0, 0]]
@@ -55,6 +60,9 @@ def test_to_dense_adj():
     assert adj.size() == (1, 10, 10)
     assert adj[0].nonzero(as_tuple=False).t().tolist() == edge_index.tolist()
 
+    adj = to_dense_adj(edge_index, batch, batch_size=4)
+    assert adj.size() == (4, 3, 3)
+
 
 def test_to_dense_adj_with_empty_edge_index():
     edge_index = torch.tensor([[], []], dtype=torch.long)

test/utils/test_to_dense_batch.py

@@ -1,35 +1,53 @@
 import torch
 
+from torch_geometric.testing import is_full_test
 from torch_geometric.utils import to_dense_batch
 
 
 def test_to_dense_batch():
     x = torch.Tensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]])
     batch = torch.tensor([0, 0, 1, 2, 2, 2])
 
-    out, mask = to_dense_batch(x, batch)
-    expected = [
+    expected = torch.Tensor([
         [[1, 2], [3, 4], [0, 0]],
         [[5, 6], [0, 0], [0, 0]],
         [[7, 8], [9, 10], [11, 12]],
-    ]
+    ])
+
+    out, mask = to_dense_batch(x, batch)
     assert out.size() == (3, 3, 2)
-    assert out.tolist() == expected
+    assert torch.equal(out, expected)
     assert mask.tolist() == [[1, 1, 0], [1, 0, 0], [1, 1, 1]]
 
+    if is_full_test():
+        jit = torch.jit.script(to_dense_batch)
+        out, mask = jit(x, batch)
+        assert torch.equal(out, expected)
+        assert mask.tolist() == [[1, 1, 0], [1, 0, 0], [1, 1, 1]]
+
+    out, mask = to_dense_batch(x, batch, max_num_nodes=2)
+    assert out.size() == (3, 2, 2)
+    assert torch.equal(out, expected[:, :2])
+    assert mask.tolist() == [[1, 1], [1, 0], [1, 1]]
+
     out, mask = to_dense_batch(x, batch, max_num_nodes=5)
     assert out.size() == (3, 5, 2)
-    assert out[:, :3].tolist() == expected
+    assert torch.equal(out[:, :3], expected)
     assert mask.tolist() == [[1, 1, 0, 0, 0], [1, 0, 0, 0, 0], [1, 1, 1, 0, 0]]
 
     out, mask = to_dense_batch(x)
     assert out.size() == (1, 6, 2)
-    assert out[0].tolist() == x.tolist()
+    assert torch.equal(out[0], x)
     assert mask.tolist() == [[1, 1, 1, 1, 1, 1]]
 
+    out, mask = to_dense_batch(x, max_num_nodes=2)
+    assert out.size() == (1, 2, 2)
+    assert torch.equal(out[0], x[:2])
+    assert mask.tolist() == [[1, 1]]
+
     out, mask = to_dense_batch(x, max_num_nodes=10)
     assert out.size() == (1, 10, 2)
-    assert out[0, :6].tolist() == x.tolist()
+    assert torch.equal(out[0, :6], x)
     assert mask.tolist() == [[1, 1, 1, 1, 1, 1, 0, 0, 0, 0]]
 
     out, mask = to_dense_batch(x, batch, batch_size=4)

torch_geometric/utils/to_dense_adj.py

@@ -12,6 +12,7 @@ def to_dense_adj(
     batch: OptTensor = None,
     edge_attr: OptTensor = None,
     max_num_nodes: Optional[int] = None,
+    batch_size: Optional[int] = None,
 ) -> Tensor:
     r"""Converts batched sparse adjacency matrices given by edge indices and
     edge attributes to a single dense batched adjacency matrix.
@@ -25,6 +26,7 @@ def to_dense_adj(
             features. (default: :obj:`None`)
         max_num_nodes (int, optional): The size of the output node dimension.
             (default: :obj:`None`)
+        batch_size (int, optional) The batch size. (default: :obj:`None`)
 
     :rtype: :class:`Tensor`
 
@@ -60,7 +62,9 @@ def to_dense_adj(
         num_nodes = int(edge_index.max()) + 1 if edge_index.numel() > 0 else 0
         batch = edge_index.new_zeros(num_nodes)
 
-    batch_size = int(batch.max()) + 1 if batch.numel() > 0 else 1
+    if batch_size is None:
+        batch_size = int(batch.max()) + 1 if batch.numel() > 0 else 1
+
     one = batch.new_ones(batch.size(0))
     num_nodes = scatter(one, batch, dim=0, dim_size=batch_size, reduce='add')
     cum_nodes = torch.cat([batch.new_zeros(1), num_nodes.cumsum(dim=0)])

torch_geometric/utils/to_dense_batch.py

@@ -98,11 +98,17 @@ def to_dense_batch(x: Tensor, batch: Optional[Tensor] = None,
                             dim_size=batch_size)
     cum_nodes = torch.cat([batch.new_zeros(1), num_nodes.cumsum(dim=0)])
 
+    filter_nodes = False
     if max_num_nodes is None:
         max_num_nodes = int(num_nodes.max())
-
-    idx = torch.arange(batch.size(0), dtype=torch.long, device=x.device)
-    idx = (idx - cum_nodes[batch]) + (batch * max_num_nodes)
+    elif num_nodes.max() > max_num_nodes:
+        filter_nodes = True
+
+    tmp = torch.arange(batch.size(0), device=x.device) - cum_nodes[batch]
+    idx = tmp + (batch * max_num_nodes)
+    if filter_nodes:
+        mask = tmp < max_num_nodes
+        x, idx = x[mask], idx[mask]
 
     size = [batch_size * max_num_nodes] + list(x.size())[1:]
     out = x.new_full(size, fill_value)