revert none return, add check

rapidsai · Jul 10, 2024 · e1fa6e0 · e1fa6e0
1 parent 4b60d8d
commit e1fa6e0
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 12 deletions.
diff --git a/python/cugraph-dgl/cugraph_dgl/tests/dataloading/test_dataloader_mg.py b/python/cugraph-dgl/cugraph_dgl/tests/dataloading/test_dataloader_mg.py
@@ -25,7 +25,7 @@
     cugraph_comms_shutdown,
 )
 
-from utils import init_pytorch_worker
+from cugraph_dgl.tests.utils import init_pytorch_worker
 
 torch = import_optional("torch")
 dgl = import_optional("dgl")

diff --git a/python/cugraph-dgl/cugraph_dgl/tests/test_graph.py b/python/cugraph-dgl/cugraph_dgl/tests/test_graph.py
@@ -52,6 +52,8 @@ def test_graph_make_homogeneous_graph(direction):
     assert (
         graph.nodes() == torch.arange(num_nodes, dtype=torch.int64, device="cuda")
     ).all()
+
+    assert graph.nodes[None]["x"] is not None
     assert (graph.nodes[None]["x"] == torch.as_tensor(node_x, device="cuda")).all()
     assert (
         graph.nodes[None]["num"]

diff --git a/python/cugraph-dgl/cugraph_dgl/tests/test_graph_mg.py b/python/cugraph-dgl/cugraph_dgl/tests/test_graph_mg.py
@@ -30,7 +30,7 @@
     cugraph_comms_get_raft_handle,
 )
 
-from utils import init_pytorch_worker
+from .utils import init_pytorch_worker
 
 pylibwholegraph = import_optional("pylibwholegraph")
 torch = import_optional("torch")
@@ -75,6 +75,7 @@ def run_test_graph_make_homogeneous_graph_mg(rank, uid, world_size, direction):
         == torch.arange(global_num_nodes, dtype=torch.int64, device="cuda")
     ).all()
     ix = torch.arange(len(node_x) * rank, len(node_x) * (rank + 1), dtype=torch.int64)
+    assert graph.nodes[ix]["x"] is not None
     assert (graph.nodes[ix]["x"] == torch.as_tensor(node_x, device="cuda")).all()
 
     assert (

diff --git a/python/cugraph-dgl/cugraph_dgl/view.py b/python/cugraph-dgl/cugraph_dgl/view.py
@@ -61,11 +61,7 @@ def __getitem__(self, key: str):
                 if self._graph._has_e_emb(t, key)
             }
 
-        return (
-            self._graph._get_e_emb(self._etype, key, self._edges)
-            if self._graph._has_e_emb(self._etype, key)
-            else None
-        )
+        return self._graph._get_e_emb(self._etype, key, self._edges)
 
     def __setitem__(self, key: str, val: Union[TensorType, Dict[str, TensorType]]):
         if isinstance(self._etype, list):
@@ -166,11 +162,7 @@ def __getitem__(self, key: str):
                 if self._graph._has_n_emb(t, key)
             }
         else:
-            return (
-                self._graph._get_n_emb(self._ntype, key, self._nodes)
-                if self._graph._has_n_emb(self._ntype, key)
-                else None
-            )
+            return self._graph._get_n_emb(self._ntype, key, self._nodes)
 
     def __setitem__(self, key: str, val: Union[TensorType, Dict[str, TensorType]]):
         if isinstance(self._ntype, list):