Cover StaticArrays

Author: simsurace

src/chainrules.jl

@@ -130,7 +130,7 @@ function ChainRulesCore.rrule(
     function pairwise_pullback(z̄)
         Δ = unthunk(z̄)
         n = size(x, dims)
-        x̄ = zero(x)
+        x̄ = collect(zero(x))
         r̄ = zero(d.r)
         if dims == 1
             for j in 1:n, i in 1:n
@@ -166,8 +166,8 @@ function ChainRulesCore.rrule(
         Δ = unthunk(z̄)
         n = size(x, dims)
         m = size(y, dims)
-        x̄ = zero(x)
-        ȳ = zero(y)
+        x̄ = collect(zero(x))
+        ȳ = collect(zero(y))
         r̄ = zero(d.r)
         if dims == 1
             for j in 1:m, i in 1:n
@@ -202,8 +202,8 @@ function ChainRulesCore.rrule(
     function colwise_pullback(z̄)
         Δ = unthunk(z̄)
         n = size(x, 2)
-        x̄ = zero(x)
-        ȳ = zero(y)
+        x̄ = collect(zero(x))
+        ȳ = collect(zero(y))
         r̄ = zero(d.r)
         for i in 1:n
             xi = view(x, :, i)

test/chainrules.jl

@@ -31,9 +31,20 @@
 
     @testset "rrules for Sinus(r=$r)" for r in (rand(3),)
         dist = KernelFunctions.Sinus(r)
-        test_rrule(dist, rand(3), rand(3))
-        test_rrule(Distances.pairwise, dist, rand(3, 2); fkwargs=(dims=2,))
-        test_rrule(Distances.pairwise, dist, rand(3, 2), rand(3, 3); fkwargs=(dims=2,))
-        test_rrule(Distances.colwise, dist, rand(3, 2), rand(3, 2))
+        @testset "$type" for type in (Vector, SVector{3})
+            test_rrule(dist, type(rand(3)), type(rand(3)))
+        end
+        @testset "$type1, $type2" for type1 in (Matrix, SMatrix{3, 2}),
+            type2 in (Matrix, SMatrix{3, 4})
+            test_rrule(
+                Distances.pairwise, dist, type1(rand(3, 2));
+                fkwargs=(dims=2,)
+            )
+            test_rrule(
+                Distances.pairwise, dist, type1(rand(3, 2)), type2(rand(3, 4));
+                fkwargs=(dims=2,)
+            )
+            test_rrule(Distances.colwise, dist, type1(rand(3, 2)), type1(rand(3, 2)))
+        end
     end
 end