Repository: pyg-team/pytorch_geometric
Pull request #7392 (2023-05-18 13:45:59)

Fix broken master


Patch:
diff --git a/torch_geometric/nn/aggr/base.py b/torch_geometric/nn/aggr/base.py
index 0dc2166cd..3de22b8e7 100644
--- a/torch_geometric/nn/aggr/base.py
+++ b/torch_geometric/nn/aggr/base.py
@@ -1,4 +1,4 @@
-from typing import Optional, Tuple, Union
+from typing import Optional, Tuple
 
 import torch
 from torch import Tensor
@@ -161,7 +161,7 @@ class Aggregation(torch.nn.Module):
         ptr: Optional[Tensor] = None,
         dim_size: Optional[int] = None,
         dim: int = -2,
-        fill_value: Union[Optional[float], Tensor] = None,
+        fill_value: float = 0.0,
         max_num_elements: Optional[int] = None,
     ) -> Tuple[Tensor, Tensor]:
 
diff --git a/torch_geometric/nn/pool/connect/base.py b/torch_geometric/nn/pool/connect/base.py
index 262937603..fd78a5c61 100644
--- a/torch_geometric/nn/pool/connect/base.py
+++ b/torch_geometric/nn/pool/connect/base.py
@@ -7,7 +7,7 @@ from torch import Tensor
 from torch_geometric.nn.pool.select import SelectOutput
 
 
-@torch.jit.script_if_tracing
+@torch.jit.script
 @dataclass(init=False)
 class ConnectOutput:
     r"""The output of the :class:`Connect` method, which holds the coarsened
diff --git a/torch_geometric/nn/pool/select/base.py b/torch_geometric/nn/pool/select/base.py
index ce9281653..fcde48234 100644
--- a/torch_geometric/nn/pool/select/base.py
+++ b/torch_geometric/nn/pool/select/base.py
@@ -5,7 +5,7 @@ import torch
 from torch import Tensor
 
 
-@torch.jit.script_if_tracing
+@torch.jit.script
 @dataclass(init=False)
 class SelectOutput:
     r"""The output of the :class:`Select` method, which holds an assignment
diff --git a/torch_geometric/utils/to_dense_batch.py b/torch_geometric/utils/to_dense_batch.py
index c90088181..1e5538fe3 100644
--- a/torch_geometric/utils/to_dense_batch.py
+++ b/torch_geometric/utils/to_dense_batch.py
@@ -1,4 +1,4 @@
-from typing import Optional, Tuple, Union
+from typing import Optional, Tuple
 
 import torch
 from torch import Tensor
@@ -6,22 +6,10 @@ from torch import Tensor
 from torch_geometric.utils import scatter
 
 
-@torch.jit._overload
-def to_dense_batch(x, batch, fill_value, max_num_nodes, batch_size):  # noqa
-    # type: (Tensor, Optional[Tensor], Optional[float], Optional[int], Optional[int]) -> Tuple[Tensor, Tensor]  # noqa
-    pass
-
-
-@torch.jit._overload
-def to_dense_batch(x, batch, fill_value, max_num_nodes, batch_size):  # noqa
-    # type: (Tensor, Optional[Tensor], Tensor, Optional[int], Optional[int]) -> Tuple[Tensor, Tensor]  # noqa
-    pass
-
-
-def to_dense_batch(  # noqa
+def to_dense_batch(
     x: Tensor,
     batch: Optional[Tensor] = None,
-    fill_value: Union[Optional[float], Tensor] = None,
+    fill_value: float = 0.0,
     max_num_nodes: Optional[int] = None,
     batch_size: Optional[int] = None,
 ) -> Tuple[Tensor, Tensor]:
@@ -41,8 +29,8 @@ def to_dense_batch(  # noqa
         batch (LongTensor, optional): Batch vector
             :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
             node to a specific example. Must be ordered. (default: :obj:`None`)
-        fill_value (float or torch.Tensor, optional): The value for invalid
-            entries in the resulting dense output tensor. (default: :obj:`0`)
+        fill_value (float, optional): The value for invalid entries in the
+            resulting dense output tensor. (default: :obj:`0`)
         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`)


Test patch:
diff --git a/test/nn/pool/select/test_select_topk.py b/test/nn/pool/select/test_select_topk.py
index 571eaa57a..a0418fff9 100644
--- a/test/nn/pool/select/test_select_topk.py
+++ b/test/nn/pool/select/test_select_topk.py
@@ -47,7 +47,7 @@ def test_select_topk(min_score):
         assert str(pool) == 'SelectTopK(16, min_score=2.0)'
 
     out = pool(x, batch)
-    assert isinstance(out, SelectOutput.__original_fn)
+    assert isinstance(out, SelectOutput)
 
     assert out.num_nodes == 6
     assert out.num_clusters <= out.num_nodes
diff --git a/test/utils/test_to_dense_batch.py b/test/utils/test_to_dense_batch.py
index 00723569e..0c84aef88 100644
--- a/test/utils/test_to_dense_batch.py
+++ b/test/utils/test_to_dense_batch.py
@@ -57,12 +57,16 @@ def test_to_dense_batch(fill):
 @onlyFullTest
 def test_to_dense_batch_jit():
     @torch.jit.script
-    def to_dense_batch_jit(x: Tensor, batch: Tensor) -> Tuple[Tensor, Tensor]:
-        return to_dense_batch(x, batch)
+    def to_dense_batch_jit(
+        x: Tensor,
+        batch: Tensor,
+        fill_value: Tensor,
+    ) -> Tuple[Tensor, Tensor]:
+        return to_dense_batch(x, batch, fill_value=fill_value)
 
     x = torch.randn(6, 2)
     batch = torch.tensor([0, 0, 1, 2, 2, 2])
 
-    out, mask = to_dense_batch_jit(x, batch)
+    out, mask = to_dense_batch_jit(x, batch, fill_value=torch.tensor(0.0))
     assert out.size() == (3, 3, 2)
     assert mask.size() == (3, 3)