faster unbatch

[CarloLucibello]

Faster implementation of unbatch for the coo type working under the assumption that the graph_indicator and the edges are sorted.

The speedup can be more than 30X and now batching and unbatching take comparable times:

using GraphNeuralNetworks, Flux
using BenchmarkTools

n = 100
c = 6
ngraphs = 128
gs = [rand_graph(n, c*n, ndata=rand(64, n), edata=rand(64, c*n)) for _ in 1:ngraphs]
gall = Flux.batch(gs)

@btime Flux.batch($gs);
# 9.858 ms (12857 allocations: 47.69 MiB)
@btime Flux.unbatch($gall);
# 335.357 ms (17801 allocations: 50.69 MiB) # master
# 10.017 ms (11807 allocations: 46.94 MiB)    # THIS PR 

[codecov[bot]]

Codecov Report

Merging #248 (65258ec) into master (2020e58) will increase coverage by 1.28%. The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master     #248      +/-   ##
==========================================
+ Coverage   80.68%   81.96%   +1.28%     
==========================================
  Files          17       17              
  Lines        1734     1846     +112     
==========================================
+ Hits         1399     1513     +114     
+ Misses        335      333       -2     

Impacted Files	Coverage Δ
src/GNNGraphs/transform.jl	97.83% <100.00%> (+1.28%)	:arrow_up:
src/layers/conv.jl	79.28% <0.00%> (+2.34%)	:arrow_up:
src/GNNGraphs/datastore.jl	76.13% <0.00%> (+3.40%)	:arrow_up:

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.