improvement for differentiability with heterogeneous graphs (#304)

* rrule for dictionary construction

* cleanup

* fix tests

* fix tests

* more tests

Author: CarloLucibello

src/GNNGraphs/GNNGraphs.jl

@@ -18,6 +18,8 @@ import MLUtils
 using MLUtils: getobs, numobs
 import Functors
 
+include("chainrules.jl") # hacks for differentiability
+
 include("datastore.jl")
 export DataStore
 

src/GNNGraphs/chainrules.jl

@@ -0,0 +1,15 @@
+# Taken from https://github.com/JuliaDiff/ChainRules.jl/pull/648
+# Remove when merged
+
+function ChainRulesCore.rrule(::Type{T}, ps::Pair...) where {T<:Dict}
+    ks = map(first, ps)
+    project_ks, project_vs = map(ProjectTo, ks), map(ProjectTo∘last, ps)
+    function Dict_pullback(ȳ)
+        dps = map(ks, project_ks, project_vs) do k, proj_k, proj_v
+            dk, dv = proj_k(getkey(ȳ, k, NoTangent())), proj_v(get(ȳ, k, NoTangent()))
+            Tangent{Pair{typeof(dk), typeof(dv)}}(first = dk, second = dv)
+        end
+       return (NoTangent(), dps...)
+    end
+    return T(ps...), Dict_pullback
+end

src/GNNGraphs/convert.jl

@@ -23,7 +23,7 @@ function to_coo(data::EDict; num_nodes = nothing, kws...)
         _num_nodes[k[1]] = max(get(_num_nodes, k[1], 0), nnodes[1])
         _num_nodes[k[3]] = max(get(_num_nodes, k[3], 0), nnodes[2])
     end
-    graph = Dict(k => v for (k, v) in pairs(graph)) # try to restrict the key/value types
+    graph = Dict([k => v for (k, v) in pairs(graph)]...) # try to restrict the key/value types
     return graph, _num_nodes, num_edges
 end
 

src/GNNGraphs/gatherscatter.jl

@@ -1,13 +1,13 @@
 _gather(x::NamedTuple, i) = map(x -> _gather(x, i), x)
-_gather(x::Dict, i) = Dict(k => _gather(v, i) for (k, v) in x)
+_gather(x::Dict, i) = Dict([k => _gather(v, i) for (k, v) in x]...)
 _gather(x::Tuple, i) = map(x -> _gather(x, i), x)
 _gather(x::AbstractArray, i) = NNlib.gather(x, i)
 _gather(x::Nothing, i) = nothing
 
 _scatter(aggr, src::Nothing, idx, n) = nothing
 _scatter(aggr, src::NamedTuple, idx, n) = map(s -> _scatter(aggr, s, idx, n), src)
 _scatter(aggr, src::Tuple, idx, n) = map(s -> _scatter(aggr, s, idx, n), src)
-_scatter(aggr, src::Dict, idx, n) = Dict(k => _scatter(aggr, v, idx, n) for (k, v) in src)
+_scatter(aggr, src::Dict, idx, n) = Dict([k => _scatter(aggr, v, idx, n) for (k, v) in src]...)
 
 function _scatter(aggr,
                   src::AbstractArray,

src/GNNGraphs/gnnheterograph.jl

@@ -99,16 +99,16 @@ GNNHeteroGraph(data; kws...) = GNNHeteroGraph(Dict(data); kws...)
 
 function GNNHeteroGraph(data::Dict; kws...)
     all(k -> k isa EType, keys(data)) || throw(ArgumentError("Keys of data must be tuples of the form (source_type, edge_type, target_type)"))
-    return GNNHeteroGraph(Dict(k => v for (k, v) in pairs(data)); kws...)
+    return GNNHeteroGraph(Dict([k => v for (k, v) in pairs(data)]...); kws...)
 end
 
 function GNNHeteroGraph(data::EDict;
                         num_nodes = nothing,
                         graph_indicator = nothing,
                         graph_type = :coo,
                         dir = :out,
-                        ndata = NDict{DataStore}(),
-                        edata = EDict{DataStore}(),
+                        ndata = nothing,
+                        edata = nothing,
                         gdata = (;))
     @assert graph_type ∈ [:coo, :dense, :sparse] "Invalid graph_type $graph_type requested"
     @assert dir ∈ [:in, :out]
@@ -132,8 +132,8 @@ function GNNHeteroGraph(data::EDict;
     num_graphs = !isnothing(graph_indicator) ?
                  maximum([maximum(gi) for gi in values(graph_indicator)]) : 1
 
-    ndata = normalize_heterographdata(ndata, default_name = :x, n = num_nodes)
-    edata = normalize_heterographdata(edata, default_name = :e, n = num_edges,
+    ndata = normalize_heterographdata(ndata, default_name = :x, ns = num_nodes)
+    edata = normalize_heterographdata(edata, default_name = :e, ns = num_edges,
                                       duplicate_if_needed = true)
     gdata = normalize_graphdata(gdata, default_name = :u, n = num_graphs)
 
@@ -226,7 +226,6 @@ num_node_types(g::GNNGraph) = 1
 
 num_node_types(g::GNNHeteroGraph) = length(g.ntypes)
 
-
 """
     edge_type_subgraph(g::GNNHeteroGraph, edge_ts)
 
@@ -240,16 +239,16 @@ function edge_type_subgraph(g::GNNHeteroGraph, edge_ts::AbstractVector{<:EType})
         @assert edge_t in g.etypes "Edge type $(edge_t) not found in graph"
     end
     node_ts = _ntypes_from_edges(edge_ts)
-    graph = Dict(edge_t => g.graph[edge_t] for edge_t in edge_ts)
-    num_nodes = Dict(node_t => g.num_nodes[node_t] for node_t in node_ts)
-    num_edges = Dict(edge_t => g.num_edges[edge_t] for edge_t in edge_ts)
+    graph = Dict([edge_t => g.graph[edge_t] for edge_t in edge_ts]...)
+    num_nodes = Dict([node_t => g.num_nodes[node_t] for node_t in node_ts]...)
+    num_edges = Dict([edge_t => g.num_edges[edge_t] for edge_t in edge_ts]...)
     if g.graph_indicator === nothing
         graph_indicator = nothing
     else
-        graph_indicator = Dict(node_t => g.graph_indicator[node_t] for node_t in node_ts)
+        graph_indicator = Dict([node_t => g.graph_indicator[node_t] for node_t in node_ts]...)
     end
-    ndata = Dict(node_t => g.ndata[node_t] for node_t in node_ts if node_t in keys(g.ndata))
-    edata = Dict(edge_t => g.edata[edge_t] for edge_t in edge_ts if edge_t in keys(g.edata))
+    ndata = Dict([node_t => g.ndata[node_t] for node_t in node_ts if node_t in keys(g.ndata)]...)
+    edata = Dict([edge_t => g.edata[edge_t] for edge_t in edge_ts if edge_t in keys(g.edata)]...)
 
     return GNNHeteroGraph(graph, num_nodes, num_edges, g.num_graphs,
                           graph_indicator, ndata, edata, g.gdata,
@@ -258,7 +257,7 @@ end
 
 # TODO this is not correct but Zygote cannot differentiate
 # through dictionary generation
-@non_differentiable edge_type_subgraph(::Any...)
+# @non_differentiable edge_type_subgraph(::Any...)
 
 function _ntypes_from_edges(edge_ts::AbstractVector{<:EType})
     ntypes = Symbol[]

src/GNNGraphs/utils.jl

@@ -163,12 +163,16 @@ function normalize_graphdata(data::NamedTuple; default_name, n, duplicate_if_nee
 end
 
 # For heterogeneous graphs
+function normalize_heterographdata(data::Nothing; default_name::Symbol, ns::Dict, kws...)
+    Dict([k => normalize_graphdata(nothing; default_name = default_name, n, kws...)
+         for (k, n) in ns]...)
+end
+
 normalize_heterographdata(data; kws...) = normalize_heterographdata(Dict(data); kws...)
 
-function normalize_heterographdata(data::Dict; default_name::Symbol, n::Dict, kws...)
-    isempty(data) && return data
-    Dict(k => normalize_graphdata(v; default_name = default_name, n = n[k], kws...)
-         for (k, v) in data)
+function normalize_heterographdata(data::Dict; default_name::Symbol, ns::Dict, kws...)
+    Dict([k => normalize_graphdata(get(data, k, nothing); default_name = default_name, n, kws...)
+         for (k, n) in ns]...)
 end
 
 ones_like(x::AbstractArray, T::Type, sz = size(x)) = fill!(similar(x, T, sz), 1)

src/layers/heteroconv.jl

@@ -12,7 +12,7 @@ function HeteroGraphConv(itr; aggr = +)
     return HeteroGraphConv(etypes, layers, aggr)
 end
 
-function (hgc::HeteroGraphConv)(g::GNNHeteroGraph, x::NamedTuple)
+function (hgc::HeteroGraphConv)(g::GNNHeteroGraph, x::Union{NamedTuple,Dict})
     function forw(l, et)
         sg = edge_type_subgraph(g, et)
         node1_t, _, node2_t = et

test/GNNGraphs/chainrules.jl

@@ -0,0 +1,24 @@
+@testset "dict constructor" begin
+    grad = gradient(1.) do x
+        d = Dict([:x => x, :y => 5]...)    
+        return sum(d[:x].^2)
+    end[1]
+
+    @test grad == 2
+
+    ## BROKEN Constructors
+    # grad = gradient(1.) do x
+    #     d = Dict([(:x => x), (:y => 5)])    
+    #     return sum(d[:x].^2)
+    # end[1]
+
+    # @test grad == 2
+
+
+    # grad = gradient(1.) do x
+    #     d = Dict([(:x => x), (:y => 5)])    
+    #     return sum(d[:x].^2)
+    # end[1]
+
+    # @test grad == 2
+end

test/GNNGraphs/gnnheterograph.jl

@@ -14,8 +14,8 @@
     @test hg.num_edges == Dict((:A, :rel1, :B) => 30, (:B, :rel2, :A) => 10)
     @test hg.graph_indicator === nothing
     @test hg.num_graphs == 1
-    @test hg.ndata == Dict()
-    @test hg.edata == Dict()
+    @test hg.ndata isa Dict{Symbol, DataStore}
+    @test hg.edata isa Dict{Tuple{Symbol, Symbol, Symbol}, DataStore}
     @test isempty(hg.gdata)
     @test sort(hg.ntypes) == [:A, :B]
     @test sort(hg.etypes) == [(:A, :rel1, :B), (:B, :rel2, :A)]

test/layers/heteroconv.jl

@@ -5,17 +5,26 @@
     model = HeteroGraphConv([(:A,:to,:B) => GraphConv(d => d), 
                             (:B,:to,:A) => GraphConv(d => d)])
 
-    x = (A = rand(Float32, d, n), B = rand(Float32, d, 2n))
+    for x in [
+                (A = rand(Float32, d, n), B = rand(Float32, d, 2n)),
+                Dict(:A => rand(Float32, d, n), :B => rand(Float32, d, 2n)) 
+             ]
+        # x = (A = rand(Float32, d, n), B = rand(Float32, d, 2n))
+        x = Dict(:A => rand(Float32, d, n), :B => rand(Float32, d, 2n)) 
+      
+        y = model(g, x)
 
-    y = model(g, x)
+        grad, dx = gradient((model, x) -> sum(model(g, x)[1]) + sum(model(g, x)[2].^2), model, x)
+        ngrad, ndx = ngradient((model, x) -> sum(model(g, x)[1]) + sum(model(g, x)[2].^2), model, x)
 
-    grad = gradient(model -> sum(model(g, x)[1]) + sum(model(g, x)[2].^2), model)[1]
-    ngrad = ngradient(model -> sum(model(g, x)[1]) + sum(model(g, x)[2].^2), model)[1]
+        @test grad.layers[1].weight1 ≈ ngrad.layers[1].weight1  rtol=1e-4
+        @test grad.layers[1].weight2 ≈ ngrad.layers[1].weight2  rtol=1e-4
+        @test grad.layers[1].bias ≈ ngrad.layers[1].bias        rtol=1e-4
+        @test grad.layers[2].weight1 ≈ ngrad.layers[2].weight1  rtol=1e-4
+        @test grad.layers[2].weight2 ≈ ngrad.layers[2].weight2  rtol=1e-4
+        @test grad.layers[2].bias ≈ ngrad.layers[2].bias        rtol=1e-4
 
-    @test grad.layers[1].weight1 ≈ ngrad.layers[1].weight1  rtol=1e-4
-    @test grad.layers[1].weight2 ≈ ngrad.layers[1].weight2  rtol=1e-4
-    @test grad.layers[1].bias ≈ ngrad.layers[1].bias        rtol=1e-4
-    @test grad.layers[2].weight1 ≈ ngrad.layers[2].weight1  rtol=1e-4
-    @test grad.layers[2].weight2 ≈ ngrad.layers[2].weight2  rtol=1e-4
-    @test grad.layers[2].bias ≈ ngrad.layers[2].bias        rtol=1e-4
+        @test dx[:A] ≈ ndx[:A]        rtol=1e-4
+        @test dx[:B] ≈ ndx[:B]        rtol=1e-4
+    end
 end