Add test

Author: aurorarossi

test/utils.jl

@@ -59,4 +59,43 @@
         @test z[:, 1:2] ≈ NNlib.softmax(e2[:, 1:2], dims = 2)
         @test z[:, 3:4] ≈ NNlib.softmax(e2[:, 3:4], dims = 2)
     end
+
+    @testset "topk_nodes" begin
+        A = [0.0297 0.8307 0.9140 0.6702 0.3346;
+             0.5901 0.3030 0.9280 0.6893 0.7997;
+             0.0880 0.6515 0.4451 0.7507 0.5297;
+             0.5171 0.6379 0.2695 0.8954 0.5197]
+        B = [0.3168 0.3174 0.5303 0.0804 0.3808;
+             0.1752 0.9105 0.5692 0.8489 0.0539;
+             0.1931 0.4954 0.3455 0.3934 0.0857;
+             0.5065 0.5182 0.5418 0.1520 0.3872]
+        C = [0.0297 0.0297 0.8307 0.9140 0.6702 0.3346;
+             0.5901 0.5901 0.3030 0.9280 0.6893 0.7997;
+             0.0880 0.0880 0.6515 0.4451 0.7507 0.5297;
+             0.5171 0.5171 0.6379 0.2695 0.8954 0.5197]
+        g1 = rand_graph(5, 6, ndata = (w = A,))
+        g2 = rand_graph(5, 6, ndata = (w = B,))
+        g3 = rand_graph(5, 6, edata = (e = C,))
+        g = Flux.batch([g1, g2])
+        output1 = topk_nodes(g1, :w, 3)
+        output2 = topk_nodes(g1, :w, 3; sortby = 5)
+        output3 = topk_edges(g3, :e, 3; sortby = 6)
+        output_batch = topk_nodes(g, :w, 3; sortby = 5)
+        correctout1 = [0.5901 0.8307 0.9280 0.8954 0.7997;
+                       0.5171 0.6515 0.9140 0.7507 0.5297;
+                       0.0880 0.6379 0.4451 0.6893 0.5197]
+        correctout2 = [0.5901 0.3030 0.9280 0.6893 0.7997;
+                       0.0880 0.6515 0.4451 0.7507 0.5297;
+                       0.5171 0.6379 0.2695 0.8954 0.5197]
+        correctout3 = [0.5901 0.5901 0.3030 0.9280 0.6893 0.7997;
+                       0.0880 0.0880 0.6515 0.4451 0.7507 0.5297;
+                       0.5171 0.5171 0.6379 0.2695 0.8954 0.5197]
+        correctout_batch = [0.5901 0.3030 0.9280 0.6893 0.7997 0.5065 0.5182 0.5418 0.1520 0.3872;
+                            0.0880 0.6515 0.4451 0.7507 0.5297 0.3168 0.3174 0.5303 0.0804 0.3808;
+                            0.5171 0.6379 0.2695 0.8954 0.5197 0.1931 0.4954 0.3455 0.3934 0.0857]
+        @test output1 == correctout1
+        @test output2 == correctout2
+        @test output3 == correctout3
+        @test output_batch == correctout_batch
+    end
 end