better inference for num_nodes

Author: CarloLucibello

GNNGraphs/src/gnngraph.jl

@@ -149,7 +149,7 @@ function GNNGraph(data::D;
 
     # don't force the shape of the data when there is only one graph
     gdata = normalize_graphdata(gdata, default_name = :u,
-                                n = num_graphs > 1 ? num_graphs : -1)
+                                n = num_graphs > 1 ? num_graphs : -1, glob=true)
 
     GNNGraph(graph,
              num_nodes, num_edges, num_graphs,
@@ -201,10 +201,10 @@ end
 
 function GNNGraph(g::GNNGraph; ndata = g.ndata, edata = g.edata, gdata = g.gdata,
                   graph_type = nothing)
-    ndata = normalize_graphdata(ndata, default_name = :x, n = g.num_nodes)
-    edata = normalize_graphdata(edata, default_name = :e, n = g.num_edges,
-                                duplicate_if_needed = true)
-    gdata = normalize_graphdata(gdata, default_name = :u, n = g.num_graphs)
+    ndata = normalize_graphdata(ndata, default_name=:x, n=g.num_nodes)
+    edata = normalize_graphdata(edata, default_name=:e, n=g.num_edges,
+                                duplicate_if_needed=true)
+    gdata = normalize_graphdata(gdata, default_name=:u, n=g.num_graphs, glob=true)
 
     if !isnothing(graph_type)
         if graph_type == :coo

GNNGraphs/src/gnnheterograph/gnnheterograph.jl

@@ -145,7 +145,7 @@ function GNNHeteroGraph(data::EDict;
         edata = normalize_heterographdata(edata, default_name = :e, ns = num_edges,
                                           duplicate_if_needed = true)
         gdata = normalize_graphdata(gdata, default_name = :u,
-                                        n = num_graphs > 1 ? num_graphs : -1)
+                                        n = num_graphs > 1 ? num_graphs : -1, glob=true)
     end
 
     return GNNHeteroGraph(graph,

GNNGraphs/src/query.jl

@@ -272,13 +272,13 @@ function CRC.rrule(::typeof(adjacency_matrix), g::G, T::DataType;
         function adjacency_matrix_pullback_weighted(Δ)
             dy = CRC.unthunk(Δ)
             s, t = edge_index(g)
-            @show dy s t
+            # @show dy s t
             #TODO using CRC.@thunk gives an error
             #TODO use gather when https://github.com/FluxML/NNlib.jl/issues/625 is fixed
             dw = zeros_like(w)
             idx = CartesianIndex.(s, t) #TODO remove when https://github.com/FluxML/NNlib.jl/issues/626 is fixed
             NNlib.gather!(dw, dy, idx)
-            @show dw
+            # @show dw
             dg = CRC.Tangent{G}(; graph = (CRC.NoTangent(), CRC.NoTangent(), dw))
             return (CRC.NoTangent(), dg, CRC.NoTangent())  
         end

GNNGraphs/src/utils.jl

@@ -129,16 +129,16 @@ function normalize_graphdata(data; default_name::Symbol, kws...)
     normalize_graphdata(NamedTuple{(default_name,)}((data,)); default_name, kws...)
 end
 
-function normalize_graphdata(data::NamedTuple; default_name, n, duplicate_if_needed = false)
+function normalize_graphdata(data::NamedTuple; default_name, n, duplicate_if_needed = false, glob=false)
     # This had to workaround two Zygote bugs with NamedTuples
     # https://github.com/FluxML/Zygote.jl/issues/1071
-    # https://github.com/FluxML/Zygote.jl/issues/1072
+    # https://github.com/FluxML/Zygote.jl/issues/1072 # TODO this is fixed
 
     if n > 1
         @assert all(x -> x isa AbstractArray, data) "Non-array features provided."
     end
 
-    if n <= 1
+    if n <= 1 && glob
         # If last array dimension is not 1, add a new dimension.
         # This is mostly useful to reshape global feature vectors
         # of size D to Dx1 matrices.

GNNGraphs/test/gnngraph.jl

@@ -53,6 +53,34 @@ end
     end
 end
 
+@testitem "Constructor: empty" setup=[GraphsTestModule] begin
+    g = GNNGraph(ndata=ones(2, 1))
+    @test g.num_nodes == 1
+    @test g.num_edges == 0
+    @test g.ndata.x == ones(2, 1)
+
+    g = GNNGraph(num_nodes=1)
+    @test g.num_nodes == 1
+    @test g.num_edges == 0
+    @test isempty(g.ndata) 
+
+    g = GNNGraph((Int[], Int[]); ndata=(;a=[1]))
+    @test g.num_nodes == 1
+    @test g.num_edges == 0
+    @test g.ndata.a == [1]
+
+    g = GNNGraph((Int[], Int[]); ndata=(;a=[1]), edata=(;b=Int[]), num_nodes=1)
+    @test g.num_nodes == 1
+    @test g.num_edges == 0
+    @test g.ndata.a == [1]
+    @test g.edata.b == Int[]
+
+    g = GNNGraph(; edata=(;b=Int[]))
+    @test g.num_nodes == 0
+    @test g.num_edges == 0
+    @test g.edata.b == Int[]
+end
+
 @testitem "symmetric graph" setup=[GraphsTestModule] tags=[:gpu] begin
     using .GraphsTestModule
     dev = gpu_device()