make coalesce compile friendly

test/utils/test_coalesce.py

@@ -1,11 +1,23 @@
 from typing import List, Optional, Tuple
 
+import pytest
 import torch
 from torch import Tensor
 
 from torch_geometric.utils import coalesce
 
 
+def test_coalesce_overflow():
+    edge_index = torch.tensor([[0, 1], [1, 0]])
+    with pytest.raises(ValueError):
+        coalesce(edge_index, num_nodes=int(1e10))
+
+    torch._dynamo.config.capture_dynamic_output_shape_ops = True
+    coalesce_compiled = torch.compile(coalesce, fullgraph=True)
+    with pytest.raises(torch._dynamo.exc.TorchRuntimeError):
+        coalesce_compiled(edge_index, num_nodes=int(1e10))
+
+
 def test_coalesce():
     edge_index = torch.tensor([[2, 1, 1, 0, 2], [1, 2, 0, 1, 1]])
     edge_attr = torch.tensor([[1], [2], [3], [4], [5]])
@@ -72,6 +84,10 @@ def wrapper3(
     ) -> Tuple[Tensor, List[Tensor]]:
         return coalesce(edge_index, edge_attr)
 
+    @torch.jit.script
+    def wrapper4(edge_index: Tensor, num_nodes: int) -> Tensor:
+        return coalesce(edge_index, num_nodes=num_nodes)
+
     edge_index = torch.tensor([[2, 1, 1, 0], [1, 2, 0, 1]])
     edge_attr = torch.tensor([[1], [2], [3], [4]])
 
@@ -91,3 +107,9 @@ def wrapper3(
     assert len(out[1]) == 2
     assert out[1][0].size() == edge_attr.size()
     assert out[1][1].size() == edge_attr.view(-1).size()
+
+    with pytest.raises(
+            torch.jit.Error,
+            match="builtins.ValueError: 'coalesce' will result in an overflow"
+    ):
+        wrapper4(edge_index, num_nodes=int(1e10))

torch_geometric/utils/_coalesce.py

@@ -131,8 +131,12 @@ def coalesce(  # noqa: F811
     num_edges = edge_index[0].size(0)
     num_nodes = maybe_num_nodes(edge_index, num_nodes)
 
-    if num_nodes * num_nodes > torch_geometric.typing.MAX_INT64:
-        raise ValueError("'coalesce' will result in an overflow")
+    MAX_NUM_NODES = torch_geometric.typing.MAX_INT64**0.5
+    if torch.jit.is_scripting():
+        if num_nodes > MAX_NUM_NODES:
+            raise ValueError("'coalesce' will result in an overflow")
+    else:
+        torch._check_value(num_nodes <= MAX_NUM_NODES)
 
     idx = edge_index[0].new_empty(num_edges + 1)
     idx[0] = -1
@@ -163,7 +167,7 @@ def coalesce(  # noqa: F811
     mask = idx[1:] > idx[:-1]
 
     # Only perform expensive merging in case there exists duplicates:
-    if mask.all():
+    if not torch_geometric.is_compiling() and mask.all():
         if edge_attr is None or isinstance(edge_attr, Tensor):
             return edge_index, edge_attr
         if isinstance(edge_attr, (list, tuple)):