Add backend checks for PyTorch only

python/dgl/dataloading/neighbor_sampler.py

@@ -149,7 +149,11 @@ def __init__(
     def sample_blocks(self, g, seed_nodes, exclude_eids=None):
         output_nodes = seed_nodes
         blocks = []
-        if F.device_type(g.device) == "cpu" and self.fused:
+        if (
+            F.device_type(g.device) == "cpu"
+            and F.backend_name == "pytorch"
+            and self.fused
+        ):
             if self.g != g:
                 self.mapping = {}
                 self.g = g

python/dgl/sampling/neighbor.py

@@ -475,8 +475,10 @@ def _sample_neighbors(
             raise DGLError(
                 "distributed training not supported in fused sampling"
             )
-        if F.device_type(g.device) != "cpu":
-            raise DGLError("Only cpu is supported in fused sampling")
+        if F.device_type(g.device) != "cpu" or F.backend_name != "pytorch":
+            raise DGLError(
+                "Only PyTorch backend and cpu is supported in fused sampling"
+            )
 
         if mapping is None:
             mapping = {}
@@ -500,7 +502,7 @@ def _sample_neighbors(
         for mapping_vector, src_nodes in zip(
             mapping[mapping_name], induced_nodes
         ):
-            mapping_vector[F.from_dgl_nd(src_nodes).type(torch.int64)] = -1
+            mapping_vector[F.from_dgl_nd(src_nodes).type(F.int64)] = -1
 
         new_ntypes = (g.ntypes, g.ntypes)
         ret = DGLBlock(subgidx, new_ntypes, g.etypes)

tests/python/common/sampling/test_sampling.py

@@ -1147,7 +1147,7 @@ def _test3():
 
 def test_sample_neighbors_noprob():
     _test_sample_neighbors(False, None, False)
-    if F._default_context_str != "gpu":
+    if F._default_context_str != "gpu" and F.backend_name == "pytorch":
         _test_sample_neighbors(False, None, True)
     # _test_sample_neighbors(True)
 
@@ -1158,7 +1158,7 @@ def test_sample_labors_noprob():
 
 def test_sample_neighbors_prob():
     _test_sample_neighbors(False, "prob", False)
-    if F._default_context_str != "gpu":
+    if F._default_context_str != "gpu" and F.backend_name == "pytorch":
         _test_sample_neighbors(False, "prob", True)
     # _test_sample_neighbors(True)
 
@@ -1169,7 +1169,7 @@ def test_sample_labors_prob():
 
 def test_sample_neighbors_outedge():
     _test_sample_neighbors_outedge(False, False)
-    if F._default_context_str != "gpu":
+    if F._default_context_str != "gpu" and F.backend_name == "pytorch":
         _test_sample_neighbors_outedge(False, True)
     # _test_sample_neighbors_outedge(True)
 
@@ -1206,8 +1206,10 @@ def test_sample_neighbors_topk_outedge():
 
 @pytest.mark.parametrize("fused", [False, True])
 def test_sample_neighbors_with_0deg(fused):
-    if fused and F._default_context_str == "gpu":
-        pytest.skip("Fused sampling doesn't support GPU.")
+    if fused and (
+        F._default_context_str == "gpu" or F.backend_name != "pytorch"
+    ):
+        pytest.skip("Fused sampling support CPU with backend PyTorch.")
     g = dgl.graph(([], []), num_nodes=5).to(F.ctx())
     sg = dgl.sampling.sample_neighbors(
         g,
@@ -1593,8 +1595,10 @@ def test_sample_neighbors_etype_sorted_homogeneous(format_, direction):
 @pytest.mark.parametrize("dtype", ["int32", "int64"])
 @pytest.mark.parametrize("fused", [False, True])
 def test_sample_neighbors_exclude_edges_heteroG(dtype, fused):
-    if fused and F._default_context_str == "gpu":
-        pytest.skip("Fused sampling doesn't support GPU.")
+    if fused and (
+        F._default_context_str == "gpu" or F.backend_name != "pytorch"
+    ):
+        pytest.skip("Fused sampling support CPU with backend PyTorch.")
     d_i_d_u_nodes = F.zerocopy_from_numpy(
         np.unique(np.random.randint(300, size=100, dtype=dtype))
     )
@@ -1753,8 +1757,10 @@ def contain_edge(g, sg, etype, u, v):
 @pytest.mark.parametrize("dtype", ["int32", "int64"])
 @pytest.mark.parametrize("fused", [False, True])
 def test_sample_neighbors_exclude_edges_homoG(dtype, fused):
-    if fused and F._default_context_str == "gpu":
-        pytest.skip("Fused sampling doesn't support GPU.")
+    if fused and (
+        F._default_context_str == "gpu" or F.backend_name != "pytorch"
+    ):
+        pytest.skip("Fused sampling support CPU with backend PyTorch.")
     u_nodes = F.zerocopy_from_numpy(
         np.unique(np.random.randint(300, size=100, dtype=dtype))
     )