Remove type constraints on Flux.batch for GNNHeteroGraph (#342)

* Remove type constraints on Flux.batch for GNNHeteroGraph

* Clean up patch

* Clean up patch

* Finalize revisions to Flux.batch patch

* Implement sugguestions

* Finalize new Flux.batch test case

* Update src/GNNGraphs/datastore.jl

* Create setindex! test for DataStore

* Update test/GNNGraphs/datastore.jl

---------

Co-authored-by: AarSeBail <>

Author: AarSeBail

src/GNNGraphs/datastore.jl

@@ -125,7 +125,7 @@ function Base.setproperty!(ds::DataStore, s::Symbol, x)
 end
 
 Base.getindex(ds::DataStore, s::Symbol) = getproperty(ds, s)
-Base.setindex!(ds::DataStore, s::Symbol, x) = setproperty!(ds, s, x)
+Base.setindex!(ds::DataStore, x, s::Symbol) = setproperty!(ds, s, x)
 
 function Base.show(io::IO, ds::DataStore)
     len = length(ds)

src/GNNGraphs/transform.jl

@@ -579,17 +579,26 @@ end
 
 
 function Flux.batch(gs::AbstractVector{<:GNNHeteroGraph})
+    function edge_index_nullable(g::GNNHeteroGraph{<:COO_T}, edge_t::EType)
+        if haskey(g.graph, edge_t)
+            g.graph[edge_t][1:2]
+        else
+            nothing
+        end
+    end
+
+    function get_edge_weight_nullable(g::GNNHeteroGraph{<:COO_T}, edge_t::EType)
+        get(g.graph, edge_t, (nothing, nothing, nothing))[3]
+    end
+
     @assert length(gs) > 0
     ntypes = union([g.ntypes for g in gs]...)
     etypes = union([g.etypes for g in gs]...)
-    # TODO remove these constraints
-    @assert ntypes == gs[1].ntypes
-    @assert etypes == gs[1].etypes
 
-    v_num_nodes = Dict(node_t => [get(g.num_nodes,node_t,0) for g in gs] for node_t in ntypes)
+    v_num_nodes = Dict(node_t => [get(g.num_nodes, node_t, 0) for g in gs] for node_t in ntypes)
     num_nodes = Dict(node_t => sum(v_num_nodes[node_t]) for node_t in ntypes)
-    num_edges = Dict(edge_t => sum(g.num_edges[edge_t] for g in gs) for edge_t in etypes)
-    edge_indices = Dict(edge_t => [edge_index(g, edge_t) for g in gs] for edge_t in etypes)
+    num_edges = Dict(edge_t => sum(get(g.num_edges, edge_t, 0) for g in gs) for edge_t in etypes)
+    edge_indices = edge_indices = Dict(edge_t => [edge_index_nullable(g, edge_t) for g in gs] for edge_t in etypes)
     nodesum = Dict(node_t => cumsum([0; v_num_nodes[node_t]])[1:(end - 1)] for node_t in ntypes)
     graphs = []
     for edge_t in etypes
@@ -599,34 +608,39 @@ function Flux.batch(gs::AbstractVector{<:GNNHeteroGraph})
         #     @show ei[1]
         # end 
         # # [ei[1] for (ii, ei) in enumerate(edge_indices[edge_t])]
-        s = cat_features([ei[1] .+ nodesum[src_t][ii] for (ii, ei) in enumerate(edge_indices[edge_t])])
-        t = cat_features([ei[2] .+ nodesum[dst_t][ii] for (ii, ei) in enumerate(edge_indices[edge_t])])
-        w = cat_features([get_edge_weight(g, edge_t) for g in gs])
+        s = cat_features([ei[1] .+ nodesum[src_t][ii] for (ii, ei) in enumerate(edge_indices[edge_t]) if ei !== nothing])
+        t = cat_features([ei[2] .+ nodesum[dst_t][ii] for (ii, ei) in enumerate(edge_indices[edge_t]) if ei !== nothing])
+        w = cat_features(filter(x -> x !== nothing, [get_edge_weight_nullable(g, edge_t) for g in gs]))
         push!(graphs, edge_t => (s, t, w))
     end
     graph = Dict(graphs...)
 
     #TODO relax this restriction
     @assert all(g -> g.num_graphs == 1, gs) 
 
-    s = edge_index(gs[1], etypes[1])[1] # grab any source vector
+    s = edge_index(gs[1], gs[1].etypes[1])[1] # grab any source vector
 
-    function materialize_graph_indicator(g, node_t)     
-        ones_like(s, g.num_nodes[node_t])
+    function materialize_graph_indicator(g, node_t)
+        n = get(g.num_nodes, node_t, 0)
+        return ones_like(s, n)
     end
     v_gi = Dict(node_t => [materialize_graph_indicator(g, node_t) for g in gs] for node_t in ntypes)
     v_num_graphs = [g.num_graphs for g in gs]
     graphsum = cumsum([0; v_num_graphs])[1:(end - 1)]
     v_gi = Dict(node_t => [ng .+ gi for (ng, gi) in zip(graphsum, v_gi[node_t])] for node_t in ntypes)
     graph_indicator = Dict(node_t => cat_features(v_gi[node_t]) for node_t in ntypes)
 
+    function data_or_else(data, types)
+        Dict(type => get(data, type, DataStore(0)) for type in types)
+    end
+
     return  GNNHeteroGraph(graph,
                 num_nodes,
                 num_edges,
                 sum(v_num_graphs),
                 graph_indicator,
-                cat_features([g.ndata for g in gs]),
-                cat_features([g.edata for g in gs]),
+                cat_features([data_or_else(g.ndata, ntypes) for g in gs]),
+                cat_features([data_or_else(g.edata, etypes) for g in gs]),
                 cat_features([g.gdata for g in gs]),
                 ntypes, etypes)
 end

src/GNNGraphs/utils.jl

@@ -119,13 +119,11 @@ end
 
 function cat_features(xs::AbstractVector{<:Dict})
     _allkeys = [sort(collect(keys(x))) for x in xs]
-    _keys = _allkeys[1]
-    all(y -> y == _keys, _allkeys) ||
-        @error "cannot concatenate feature data with different keys"
+    _keys = union(_allkeys...)
     length(xs) == 1 && return xs[1]
 
     # concatenate 
-    return Dict([k => cat_features([x[k] for x in xs]) for k in _keys]...)
+    return Dict([k => cat_features([x[k] for x in xs if haskey(x, k)]) for k in _keys]...)
 end
 
 

test/GNNGraphs/datastore.jl

@@ -27,6 +27,13 @@ end
     @test vec._data == [Dict(:x => x), Dict(:x => x, :y => vec[2].y)]
 end
 
+@testset "setindex!" begin
+    ds = DataStore(10)
+    x = rand(10)
+    @test (ds[:x] = x) == x # Tests setindex!
+    @test ds.x == ds[:x] == x
+end
+
 @testset "map" begin
     ds = DataStore(10, (:x => rand(10), :y => rand(2, 10)))
     ds2 = map(x -> x .+ 1, ds)

test/GNNGraphs/transform.jl

@@ -346,7 +346,75 @@ end
         @test g.ndata[:A].x == ones(2, 50)
         @test g.ndata[:A].y == zeros(50)
         @test g.edata[(:A,:to,:B)].e == fill(2, 100)
-        @test g.gdata.u == fill(7, 5)        
+        @test g.gdata.u == fill(7, 5)
+    end
+
+    @testset "batch non-similar edge types" begin
+        gs = [rand_heterograph((:A =>10, :B => 14), ((:A, :to1, :A) => 5, (:A, :to1, :B) => 20)),
+            rand_heterograph((:A => 10, :B => 15), ((:A, :to1, :B) => 5, (:B, :to2, :B) => 16)),
+            rand_heterograph((:B => 15, :C => 5), ((:C, :to1, :B) => 5, (:B, :to2, :C) => 21)),
+            rand_heterograph((:A => 10, :B => 10, :C => 10), ((:A, :to1, :C) => 5, (:A, :to1, :B) => 5)),
+            rand_heterograph((:C => 20), ((:C, :to3, :C) => 10))
+        ]
+        g = Flux.batch(gs)
+
+        @test g.num_nodes[:A] == 10 + 10 + 10
+        @test g.num_nodes[:B] == 14 + 15 + 15 + 10
+        @test g.num_nodes[:C] == 5 + 10 + 20
+        @test g.num_edges[(:A,:to1,:A)] == 5
+        @test g.num_edges[(:A,:to1,:B)] == 20 + 5 + 5
+        @test g.num_edges[(:A,:to1,:C)] == 5
+
+        @test g.num_edges[(:B,:to2,:B)] == 16
+        @test g.num_edges[(:B,:to2,:C)] == 21
+
+        @test g.num_edges[(:C,:to1,:B)] == 5
+        @test g.num_edges[(:C,:to3,:C)] == 10
+        @test length(keys(g.num_edges)) == 7
+        @test g.num_graphs == 5
+
+        function ndata_if_key(g, key, subkey, value)
+            if haskey(g.ndata, key)
+                g.ndata[key][subkey] = reduce(hcat, fill(value, g.num_nodes[key]))
+            end
+        end
+
+        function edata_if_key(g, key, subkey, value)
+            if haskey(g.edata, key)
+                g.edata[key][subkey] = reduce(hcat, fill(value, g.num_edges[key]))
+            end
+        end
+
+        for gi in gs
+            ndata_if_key(gi, :A, :x, [0])
+            ndata_if_key(gi, :A, :y, ones(2))
+            ndata_if_key(gi, :B, :x, ones(3))
+            ndata_if_key(gi, :C, :y, zeros(4))
+            edata_if_key(gi, (:A,:to1,:B), :x, [0])
+            gi.gdata.u = 7
+        end
+
+        g = Flux.batch(gs)
+
+        @test g.ndata[:A].x == reduce(hcat, fill(0, 10 + 10 + 10))
+        @test g.ndata[:A].y == ones(2, 10 + 10 + 10)
+        @test g.ndata[:B].x == ones(3, 14 + 15 + 15 + 10)
+        @test g.ndata[:C].y == zeros(4, 5 + 10 + 20)
+
+        @test g.edata[(:A,:to1,:B)].x == reduce(hcat, fill(0, 20 + 5 + 5))
+
+        @test g.gdata.u == fill(7, 5)
+
+        # Allow for wider eltype 
+        g = Flux.batch(GNNHeteroGraph[g for g in gs])
+        @test g.ndata[:A].x == reduce(hcat, fill(0, 10 + 10 + 10))
+        @test g.ndata[:A].y == ones(2, 10 + 10 + 10)
+        @test g.ndata[:B].x == ones(3, 14 + 15 + 15 + 10)
+        @test g.ndata[:C].y == zeros(4, 5 + 10 + 20)
+
+        @test g.edata[(:A,:to1,:B)].x == reduce(hcat, fill(0, 20 + 5 + 5))
+
+        @test g.gdata.u == fill(7, 5)
     end
 
     @testset "add_edges" begin