The constraint in Flux.batch(gs::AbstractVector{<:GNNHeteroGraph}) does not seem to be strong enough

[AarSeBail]

While I'd also like there to be no constraint, I do not believe that the current constraint is sufficient.

https://github.com/CarloLucibello/GraphNeuralNetworks.jl/blob/066d8f50c88f4e558be50c9f45b7fbdea8bb6527/src/GNNGraphs/transform.jl#L585C3-L585C3

Consider for example the following:

using GraphNeuralNetworks
using Flux

gs = [rand_heterograph((:A =>10, :B => 14), ((:A, :to1, :A) => 5, (:A, :to2, :B) => 20)),
                   rand_heterograph((:A =>10, :B => 14), ((:A, :to1, :A) => 5))]

g = Flux.batch(gs)

The code errors with ERROR: KeyError: key (:A, :to2, :B) not found due to

num_edges = Dict(edge_t => sum(g.num_edges[edge_t] for g in gs) for edge_t in etypes)

I believe that the assertion needs to be something along the following lines.

# TODO remove these constraints
@assert ntypes == intersect([g.ntypes for g in gs]...)
@assert etypes == intersect([g.etypes for g in gs]...)