add_edges for GNNHeteroGraph does not allow providing the number of nodes

[svilupp]

I think I noticed a bug when adding new edges (with new nodes) into GNNHeteroGraph.

Documentation for add_edges says:

If the edge type is not already present in the graph, it is added. If it involves new node types, they are added to the graph as well. In this case, a dictionary or named tuple of num_nodes can be passed to specify the number of nodes of the new types, otherwise the number of nodes is inferred from the maximum

Expectation: When adding new edges with a new node, I want to provide the number of the new nodes via num_nodes kwarg

Behavior: When I provide num_nodes, I get

ERROR: KeyError: key :actor not found
Stacktrace:
[1] getindex
@ GraphNeuralNetworks.GNNGraphs ./dict.jl:496 [inlined]
[2] add_edges(g::GNNHeteroGraph{…}, edge_t::Tuple{…}, snew::Vector{…}, tnew::Vector{…}; edata::Nothing, num_nodes::Dict{…})
@ GraphNeuralNetworks.GNNGraphs ~/Documents/GitHub/GraphNeuralNetworks.jl/src/GNNGraphs/transform.jl:195
[3] add_edges
@ GraphNeuralNetworks.GNNGraphs ~/Documents/GitHub/GraphNeuralNetworks.jl/src/GNNGraphs/transform.jl:162 [inlined]
[4] #add_edges#178
@ GraphNeuralNetworks.GNNGraphs ~/Documents/GitHub/GraphNeuralNetworks.jl/src/GNNGraphs/transform.jl:160 [inlined]
[5] top-level scope

MWE:

using GraphNeuralNetworks

g = GNNHeteroGraph((:user, :rate, :movie) => ([1,1,2,3], [7,13,5,7]); num_nodes=Dict(:user=>10,:movie=>15))
g = add_edges(g, (:actor, :like, :movie) => ([1,2,3,3,3], [3,5,1,9,4]); num_nodes=Dict(:actor=>10))
# ERROR: KeyError: key :actor not found

# NamedTuple fails too
g = add_edges(g, (:actor, :like, :movie) => ([1,2,3,3,3], [3,5,1,9,4]);num_nodes=(;:actor=>10,))
# ERROR: KeyError: key :actor not found

# this works:
g = add_edges(g, (:actor, :like, :movie) => ([1,2,3,3,3], [3,5,1,9,4]))

Versioninfo:

julia> versioninfo()
Julia Version 1.10.0-beta2
Commit a468aa198d0 (2023-08-17 06:27 UTC)
Build Info:
Official https://julialang.org/ release
Platform Info:
OS: macOS (arm64-apple-darwin22.4.0)
CPU: 8 × Apple M1 Pro
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-15.0.7 (ORCJIT, apple-m1)
Threads: 8 on 6 virtual cores
Environment:
JULIA_EDITOR = code
JULIA_NUM_THREADS = 8

Package version: 0.6.11 (current main branch)

[svilupp]

I'll open a PR with a fix later.

[CarloLucibello]

Thanks, a PR would be appreciated. Heterographs support is still quite rough and needs more love.