Heterogeneous Link Prediction Example for cuGraph-PyG

conda/environments/all_cuda-118_arch-x86_64.yaml

@@ -39,6 +39,7 @@ dependencies:
 - pytest-forked
 - pytest-xdist
 - pytorch>=2.3
+- pytorch>=2.3,<2.6a0
 - pytorch_geometric>=2.5,<2.7
 - rapids-build-backend>=0.3.0,<0.4.0.dev0
 - rmm==25.2.*,>=0.0.0a0

conda/environments/all_cuda-121_arch-x86_64.yaml

@@ -45,6 +45,7 @@ dependencies:
 - pytest-forked
 - pytest-xdist
 - pytorch>=2.3
+- pytorch>=2.3,<2.6a0
 - pytorch_geometric>=2.5,<2.7
 - rapids-build-backend>=0.3.0,<0.4.0.dev0
 - rmm==25.2.*,>=0.0.0a0

conda/environments/all_cuda-124_arch-x86_64.yaml

@@ -45,6 +45,7 @@ dependencies:
 - pytest-forked
 - pytest-xdist
 - pytorch>=2.3
+- pytorch>=2.3,<2.6a0
 - pytorch_geometric>=2.5,<2.7
 - rapids-build-backend>=0.3.0,<0.4.0.dev0
 - rmm==25.2.*,>=0.0.0a0

dependencies.yaml

@@ -20,6 +20,7 @@ files:
       - depends_on_dask_cudf
       - depends_on_cupy
       - depends_on_pytorch
+      - depends_on_ogb
       - depends_on_dgl
       - depends_on_pyg
       - python_run_cugraph_dgl
@@ -52,6 +53,7 @@ files:
     includes:
       - cuda_version
       - depends_on_pytorch
+      - depends_on_ogb
       - depends_on_cugraph_dgl
       - py_version
       - test_notebook
@@ -63,6 +65,7 @@ files:
       - depends_on_cudf
       - depends_on_dgl
       - depends_on_pytorch
+      - depends_on_ogb
       - py_version
       - test_python_common
 
@@ -123,6 +126,7 @@ files:
     includes:
       - depends_on_pylibwholegraph
       - depends_on_pytorch
+      - depends_on_ogb
       - test_python_common
   py_build_cugraph_pyg:
     output: pyproject
@@ -150,6 +154,7 @@ files:
       - depends_on_pyg
       - depends_on_pylibwholegraph
       - depends_on_pytorch
+      - depends_on_ogb
       - test_python_common
       - test_python_cugraph_pyg
 
@@ -163,6 +168,7 @@ files:
       - depends_on_cugraph
       - depends_on_dgl
       - depends_on_pytorch
+      - depends_on_ogb
       - cugraph_dgl_dev
       - test_python_common
   cugraph_pyg_dev:
@@ -175,6 +181,7 @@ files:
       - depends_on_cugraph
       - depends_on_pyg
       - depends_on_pytorch
+      - depends_on_ogb
       - cugraph_pyg_dev
       - test_python_common
 channels:
@@ -428,6 +435,34 @@ dependencies:
               - *tensordict
           - {matrix: null, packages: [*pytorch_pip, *tensordict]}
 
+  # Will remove this after snap-stanford/ogb#497 is resolved.
+  # Temporarily sets the max pytorch version to 2.5 for compatibility
+  # with ogb.
+  depends_on_ogb:
+    common:
+      - output_types: [conda]
+        packages:
+          - pytorch>=2.3,<2.6a0
+    specific:
+      - output_types: [requirements]
+        matrices:
+          - matrix: {cuda: "12.*"}
+            packages:
+              - --extra-index-url=https://download.pytorch.org/whl/cu121
+          - matrix: {cuda: "11.*"}
+            packages:
+              - --extra-index-url=https://download.pytorch.org/whl/cu118
+          - {matrix: null, packages: null}
+      - output_types: [requirements, pyproject]
+        matrices:
+          - matrix: {cuda: "12.*"}
+            packages:
+              - torch>=2.3,<2.6a0
+          - matrix: {cuda: "11.*"}
+            packages:
+              - torch>=2.3,<2.6a0
+          - {matrix: null, packages: [*pytorch_pip]}
+
   depends_on_dgl:
     specific:
       - output_types: [conda]

python/cugraph-dgl/conda/cugraph_dgl_dev_cuda-118.yaml

@@ -17,6 +17,7 @@ dependencies:
 - pytest-cov
 - pytest-xdist
 - pytorch>=2.3
+- pytorch>=2.3,<2.6a0
 - tensordict>=0.1.2
 - torchdata
 name: cugraph_dgl_dev_cuda-118

python/cugraph-pyg/conda/cugraph_pyg_dev_cuda-118.yaml

@@ -15,6 +15,7 @@ dependencies:
 - pytest-cov
 - pytest-xdist
 - pytorch>=2.3
+- pytorch>=2.3,<2.6a0
 - pytorch_geometric>=2.5,<2.7
 - tensordict>=0.1.2
 - torchdata

python/cugraph-pyg/cugraph_pyg/data/graph_store.py

@@ -293,7 +293,9 @@ def __get_weight_tensor(
         return torch.concat(weights)
 
     @property
-    def _numeric_edge_types(self) -> Tuple[List, "torch.Tensor", "torch.Tensor"]:
+    def _numeric_edge_types(
+        self,
+    ) -> Tuple[List[Tuple[str, str, str]], "torch.Tensor", "torch.Tensor"]:
         """
         Returns the canonical edge types in order (the 0th canonical type corresponds
         to numeric edge type 0, etc.), along with the numeric source and destination