Move torch.jit.script check to test (#194)

* update

* update

* update

Author: rusty1s

CMakeLists.txt

@@ -1,7 +1,7 @@
 cmake_minimum_required(VERSION 3.0)
 project(torchcluster)
 set(CMAKE_CXX_STANDARD 14)
-set(TORCHCLUSTER_VERSION 1.6.2)
+set(TORCHCLUSTER_VERSION 1.6.3)
 
 option(WITH_CUDA "Enable CUDA support" OFF)
 option(WITH_PYTHON "Link to Python when building" ON)

conda/pytorch-cluster/meta.yaml

@@ -1,6 +1,6 @@
 package:
   name: pytorch-cluster
-  version: 1.6.2
+  version: 1.6.3
 
 source:
   path: ../..

setup.py

@@ -11,7 +11,7 @@
 from torch.utils.cpp_extension import (CUDA_HOME, BuildExtension, CppExtension,
                                        CUDAExtension)
 
-__version__ = '1.6.2'
+__version__ = '1.6.3'
 URL = 'https://github.com/rusty1s/pytorch_cluster'
 
 WITH_CUDA = False

test/test_graclus.py

@@ -50,3 +50,7 @@ def test_graclus_cluster(test, dtype, device):
 
     cluster = graclus_cluster(row, col, weight)
     assert_correct(row, col, cluster)
+
+    jit = torch.jit.script(graclus_cluster)
+    cluster = jit(row, col, weight)
+    assert_correct(row, col, cluster)

test/test_grid.py

@@ -38,3 +38,6 @@ def test_grid_cluster(test, dtype, device):
 
     cluster = grid_cluster(pos, size, start, end)
     assert cluster.tolist() == test['cluster']
+
+    jit = torch.jit.script(grid_cluster)
+    assert torch.equal(jit(pos, size, start, end), cluster)

test/test_knn.py

@@ -34,6 +34,10 @@ def test_knn(dtype, device):
     edge_index = knn(x, y, 2)
     assert to_set(edge_index) == set([(0, 2), (0, 3), (1, 0), (1, 1)])
 
+    jit = torch.jit.script(knn)
+    edge_index = jit(x, y, 2)
+    assert to_set(edge_index) == set([(0, 2), (0, 3), (1, 0), (1, 1)])
+
     edge_index = knn(x, y, 2, batch_x, batch_y)
     assert to_set(edge_index) == set([(0, 2), (0, 3), (1, 4), (1, 5)])
 
@@ -65,6 +69,11 @@ def test_knn_graph(dtype, device):
     assert to_set(edge_index) == set([(1, 0), (3, 0), (0, 1), (2, 1), (1, 2),
                                       (3, 2), (0, 3), (2, 3)])
 
+    jit = torch.jit.script(knn_graph)
+    edge_index = jit(x, k=2, flow='source_to_target')
+    assert to_set(edge_index) == set([(1, 0), (3, 0), (0, 1), (2, 1), (1, 2),
+                                      (3, 2), (0, 3), (2, 3)])
+
 
 @pytest.mark.parametrize('dtype,device', product([torch.float], devices))
 def test_knn_graph_large(dtype, device):

test/test_radius.py

@@ -35,6 +35,11 @@ def test_radius(dtype, device):
     assert to_set(edge_index) == set([(0, 0), (0, 1), (0, 2), (0, 3), (1, 1),
                                       (1, 2), (1, 5), (1, 6)])
 
+    jit = torch.jit.script(radius)
+    edge_index = jit(x, y, 2, max_num_neighbors=4)
+    assert to_set(edge_index) == set([(0, 0), (0, 1), (0, 2), (0, 3), (1, 1),
+                                      (1, 2), (1, 5), (1, 6)])
+
     edge_index = radius(x, y, 2, batch_x, batch_y, max_num_neighbors=4)
     assert to_set(edge_index) == set([(0, 0), (0, 1), (0, 2), (0, 3), (1, 5),
                                       (1, 6)])
@@ -64,12 +69,20 @@ def test_radius_graph(dtype, device):
     assert to_set(edge_index) == set([(1, 0), (3, 0), (0, 1), (2, 1), (1, 2),
                                       (3, 2), (0, 3), (2, 3)])
 
+    jit = torch.jit.script(radius_graph)
+    edge_index = jit(x, r=2.5, flow='source_to_target')
+    assert to_set(edge_index) == set([(1, 0), (3, 0), (0, 1), (2, 1), (1, 2),
+                                      (3, 2), (0, 3), (2, 3)])
+
 
 @pytest.mark.parametrize('dtype,device', product([torch.float], devices))
 def test_radius_graph_large(dtype, device):
     x = torch.randn(1000, 3, dtype=dtype, device=device)
 
-    edge_index = radius_graph(x, r=0.5, flow='target_to_source', loop=True,
+    edge_index = radius_graph(x,
+                              r=0.5,
+                              flow='target_to_source',
+                              loop=True,
                               max_num_neighbors=2000)
 
     tree = scipy.spatial.cKDTree(x.cpu().numpy())

test/test_rw.py

@@ -31,6 +31,9 @@ def test_rw_small(device):
     out = random_walk(row, col, start, walk_length, num_nodes=3)
     assert out.tolist() == [[0, 1, 0, 1, 0], [1, 0, 1, 0, 1], [2, 2, 2, 2, 2]]
 
+    jit = torch.jit.script(random_walk)
+    assert torch.equal(jit(row, col, start, walk_length, num_nodes=3), out)
+
 
 @pytest.mark.parametrize('device', devices)
 def test_rw_large_with_edge_indices(device):

torch_cluster/__init__.py

@@ -3,7 +3,7 @@
 
 import torch
 
-__version__ = '1.6.2'
+__version__ = '1.6.3'
 
 for library in [
         '_version', '_grid', '_graclus', '_fps', '_rw', '_sampler', '_nearest',

torch_cluster/graclus.py

@@ -3,10 +3,12 @@
 import torch
 
 
-@torch.jit.script
-def graclus_cluster(row: torch.Tensor, col: torch.Tensor,
-                    weight: Optional[torch.Tensor] = None,
-                    num_nodes: Optional[int] = None) -> torch.Tensor:
+def graclus_cluster(
+    row: torch.Tensor,
+    col: torch.Tensor,
+    weight: Optional[torch.Tensor] = None,
+    num_nodes: Optional[int] = None,
+) -> torch.Tensor:
     """A greedy clustering algorithm of picking an unmarked vertex and matching
     it with one its unmarked neighbors (that maximizes its edge weight).