Skip to content

Commit f0f715b

Browse files
committed
Add tests for coalesced COO graphs
1 parent 7274e04 commit f0f715b

File tree

2 files changed

+14
-2
lines changed

2 files changed

+14
-2
lines changed

GNNlib/test/test_module.jl

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,7 @@ function test_gradients(
150150
return true
151151
end
152152

153-
function generate_test_graphs(graph_type)
153+
function generate_test_graphs(graph_type; do_coalesce=false)
154154
adj1 = [0 1 0 1
155155
1 0 1 0
156156
0 1 0 1
@@ -168,12 +168,18 @@ function generate_test_graphs(graph_type)
168168
g_single_vertex = GNNGraph(adj_single_vertex,
169169
ndata = rand(Float32, D_IN, 4);
170170
graph_type)
171+
172+
if graph_type == :coo && do_coalesce
173+
g1 = coalesce(g1)
174+
g_single_vertex = coalesce(g_single_vertex)
175+
end
171176

172177
return (g1, g_single_vertex)
173178
end
174179

175180
GRAPH_TYPES = [:coo, :dense, :sparse]
176181
TEST_GRAPHS = [generate_test_graphs(:coo)...,
182+
generate_test_graphs(:coo, do_coalesce=true)...,
177183
generate_test_graphs(:dense)...,
178184
generate_test_graphs(:sparse)...]
179185

GraphNeuralNetworks/test/test_module.jl

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,7 @@ function test_gradients(
157157
end
158158

159159

160-
function generate_test_graphs(graph_type)
160+
function generate_test_graphs(graph_type; do_coalesce=false)
161161
adj1 = [0 1 0 1
162162
1 0 1 0
163163
0 1 0 1
@@ -175,12 +175,18 @@ function generate_test_graphs(graph_type)
175175
g_single_vertex = GNNGraph(adj_single_vertex,
176176
ndata = rand(Float32, D_IN, 4);
177177
graph_type)
178+
179+
if graph_type == :coo && do_coalesce
180+
g1 = coalesce(g1)
181+
g_single_vertex = coalesce(g_single_vertex)
182+
end
178183

179184
return (g1, g_single_vertex)
180185
end
181186

182187
GRAPH_TYPES = [:coo, :dense, :sparse]
183188
TEST_GRAPHS = [generate_test_graphs(:coo)...,
189+
generate_test_graphs(:coo, do_coalesce=true)...,
184190
generate_test_graphs(:dense)...,
185191
generate_test_graphs(:sparse)...]
186192

0 commit comments

Comments
 (0)