JuliaGraphs · CarloLucibello · Dec 28, 2024 · Dec 26, 2024 · Dec 27, 2024 · Dec 27, 2024
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -100,3 +100,54 @@ steps:
       GNN_TEST_AMDGPU: "true"
       GNN_TEST_CPU: "false"
     timeout_in_minutes: 60
+
+  - label: "GNNGraphs CUDA"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1"
+      - JuliaCI/julia-coverage#v1:
+          dirs:
+            - GNNGraphs/src
+    command: |
+      julia --color=yes --depwarn=yes --project=GNNGraphs/test -e '
+        import Pkg
+        dev_pkgs = Pkg.PackageSpec[]
+        for pkg in ("GNNGraphs",)
+          push!(dev_pkgs, Pkg.PackageSpec(path=pkg));
+        end
+        Pkg.develop(dev_pkgs)
+        Pkg.add(["CUDA", "cuDNN"])
+        Pkg.test("GNNGraphs")'
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    env:
+      GNN_TEST_CUDA: "true"
+      GNN_TEST_CPU: "false"
+    timeout_in_minutes: 60
+
+  - label: "GNNGraphs AMDGPU"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1"
+      - JuliaCI/julia-coverage#v1:
+          dirs:
+            - GNNGraphs/src
+    command: |
+      julia --color=yes --depwarn=yes --project=GNNGraphs/test -e '
+        import Pkg
+        dev_pkgs = Pkg.PackageSpec[]
+        for pkg in ("GNNGraphs",)
+          push!(dev_pkgs, Pkg.PackageSpec(path=pkg));
+        end
+        Pkg.develop(dev_pkgs)
+        Pkg.add(["AMDGPU"])
+        Pkg.test("GNNGraphs")'
+    agents:
+      queue: "juliagpu"
+      rocm: "*"
+      rocmgpu: "*"
+    env:
+      GNN_TEST_AMDGPU: "true"
+      GNN_TEST_CPU: "false"
+    timeout_in_minutes: 60
diff --git a/.github/workflows/test_GNNGraphs.yml b/.github/workflows/test_GNNGraphs.yml
@@ -8,7 +8,7 @@ on:
       - master
 jobs:
   test:
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false

diff --git a/.github/workflows/test_GNNLux.yml b/.github/workflows/test_GNNLux.yml
@@ -8,7 +8,7 @@ on:
       - master
 jobs:
   test:
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false

diff --git a/.github/workflows/test_GNNlib.yml b/.github/workflows/test_GNNlib.yml
@@ -8,7 +8,7 @@ on:
       - master
 jobs:
   test:
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false

diff --git a/.github/workflows/test_GraphNeuralNetworks.yml b/.github/workflows/test_GraphNeuralNetworks.yml
@@ -8,7 +8,7 @@ on:
       - master
 jobs:
   test:
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false

diff --git a/GNNGraphs/Project.toml b/GNNGraphs/Project.toml
@@ -4,13 +4,13 @@ authors = ["Carlo Lucibello and contributors"]
 version = "1.4.1"
 
 [deps]
-Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
+MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 NearestNeighbors = "b8a86587-4115-5ab1-83bc-aa920d37bbce"
@@ -28,15 +28,14 @@ GNNGraphsCUDAExt = "CUDA"
 GNNGraphsSimpleWeightedGraphsExt = "SimpleWeightedGraphs"
 
 [compat]
-Adapt = "4"
 CUDA = "5"
 ChainRulesCore = "1"
 Functors = "0.5"
 Graphs = "1.4"
 KrylovKit = "0.8"
 LinearAlgebra = "1"
 MLDataDevices = "1.0"
-MLDatasets = "0.7"
+MLDatasets = "0.7.18"
 MLUtils = "0.4"
 NNlib = "0.9"
 NearestNeighbors = "0.4"
@@ -45,21 +44,4 @@ SimpleWeightedGraphs = "1.4.0"
 SparseArrays = "1"
 Statistics = "1"
 StatsBase = "0.34"
-cuDNN = "1"
 julia = "1.10"
-
-[extras]
-Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
-DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
-FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
-InlineStrings = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48"
-MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
-SimpleWeightedGraphs = "47aef6b3-ad0c-573a-a1e2-d07658019622"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
-cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
-
-[targets]
-test = ["Test", "Adapt", "DataFrames", "InlineStrings", "SimpleWeightedGraphs", "Zygote", "FiniteDifferences", "ChainRulesTestUtils", "MLDatasets", "CUDA", "cuDNN"]
diff --git a/GNNGraphs/docs/make.jl b/GNNGraphs/docs/make.jl
@@ -45,7 +45,6 @@ makedocs(;
             "GNNGraph" => "api/gnngraph.md",
             "GNNHeteroGraph" => "api/heterograph.md",
             "TemporalSnapshotsGNNGraph" => "api/temporalgraph.md",
-            "Samplers" => "api/samplers.md",
             "Datasets" => "api/datasets.md",
         ],
     ]

diff --git a/GNNGraphs/docs/src/api/gnngraph.md b/GNNGraphs/docs/src/api/gnngraph.md
@@ -77,12 +77,8 @@ Base.intersect
 
 ## Sampling 
 
-```@autodocs; canonical = true
-Modules = [GNNGraphs]
-Pages   = ["src/sampling.jl"]
-Private = false
-```
-
-```@docs; canonical = true
+```@docs
+NeighborLoader
+sample_neighbors
 Graphs.induced_subgraph(::GNNGraph, ::Vector{Int})
-```
+```
diff --git a/GNNGraphs/docs/src/api/samplers.md b/GNNGraphs/docs/src/api/samplers.md
diff --git a/GNNGraphs/src/GNNGraphs.jl b/GNNGraphs/src/GNNGraphs.jl
@@ -13,6 +13,7 @@ using LinearAlgebra, Random, Statistics
 import MLUtils
 using MLUtils: getobs, numobs, ones_like, zeros_like, chunk, batch, rand_like
 using MLDataDevices: get_device, cpu_device, CPUDevice
+using Functors: @functor
 
 include("chainrules.jl") # hacks for differentiability
 
@@ -54,7 +55,7 @@ export adjacency_list,
        normalized_laplacian,
        scaled_laplacian,
        laplacian_lambda_max,
-# from Graphs
+# from Graphs.jl
        adjacency_matrix,
        degree,
        has_edge, 
@@ -82,7 +83,7 @@ export add_nodes,
        perturb_edges,
        remove_nodes,
        ppr_diffusion,
-# from MLUtils
+# from MLUtils.jl
        batch,
        unbatch,
 # from SparseArrays
@@ -98,9 +99,6 @@ export rand_graph,
        rand_temporal_radius_graph,
        rand_temporal_hyperbolic_graph
 
-include("sampling.jl")
-export sample_neighbors
-
 include("operators.jl")
 # Base.intersect
 
@@ -117,7 +115,8 @@ export mldataset2gnngraph
 
 include("deprecations.jl")
 
-include("samplers.jl")
-export NeighborLoader
+include("sampling.jl")
+export NeighborLoader, sample_neighbors, 
+      induced_subgraph # from Graphs.jl
 
 end #module
diff --git a/GNNGraphs/src/samplers.jl b/GNNGraphs/src/samplers.jl