implement StatsBase.nobs instead of LearnBase.nobs (#62)

* implement StatsBase.nobs

* fix test

* fix test

Author: CarloLucibello

Project.toml

@@ -19,6 +19,7 @@ NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 TestEnv = "1e6cf692-eddd-4d53-88a5-2d735e33781b"
 
 [compat]
@@ -33,8 +34,9 @@ LearnBase = "0.4, 0.5"
 MacroTools = "0.5"
 NNlib = "0.7"
 NNlibCUDA = "0.1"
-julia = "1.6"
+StatsBase = "0.32, 0.33"
 TestEnv = "1"
+julia = "1.6"
 
 [extras]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/GraphNeuralNetworks.jl

@@ -10,6 +10,7 @@ using Flux
 using Flux: glorot_uniform, leakyrelu, GRUCell, @functor, batch
 using MacroTools: @forward
 import LearnBase
+import StatsBase
 using LearnBase: getobs
 using NNlib, NNlibCUDA
 using NNlib: scatter, gather

src/gnngraph.jl

@@ -499,8 +499,8 @@ Equivalent to [`SparseArrays.blockdiag`](@ref).
 """
 Flux.batch(xs::Vector{<:GNNGraph}) = blockdiag(xs...)
 
-### LearnBase compatibility
-LearnBase.nobs(g::GNNGraph) = g.num_graphs 
+### StatsBase/LearnBase compatibility
+StatsBase.nobs(g::GNNGraph) = g.num_graphs 
 LearnBase.getobs(g::GNNGraph, i) = getgraph(g, i)
 
 # Flux's Dataloader compatibility. Related PR https://github.com/FluxML/Flux.jl/pull/1683

test/gnngraph.jl

@@ -272,7 +272,7 @@
 
         @test LearnBase.getobs(g, 3) == getgraph(g, 3)
         @test LearnBase.getobs(g, 3:5) == getgraph(g, 3:5)
-        @test LearnBase.nobs(g) == g.num_graphs
+        @test StatsBase.nobs(g) == g.num_graphs
 
         d = Flux.Data.DataLoader(g, batchsize = 2, shuffle=false)
         @test first(d) == getgraph(g, 1:2)

test/runtests.jl

@@ -5,6 +5,7 @@ using Flux: gpu, @functor
 using LinearAlgebra, Statistics, Random
 using NNlib
 using LearnBase
+import StatsBase
 using Graphs
 using Zygote
 using Test