Update rrules

Author: nickrobinson251

src/chainrules.jl

@@ -27,7 +27,7 @@ function ChainRulesCore.rrule(
         ::typeof(*),
         bm::BlockDiagonal{T, V},
         v::StridedVector{T}
-    ) where {T<:Union{Real, Complex}, V<:Matrix{T}}
+    ) where {T<:Union{Real, Complex}, V}
 
     y = bm * v
 

test/chainrules.jl

@@ -1,17 +1,19 @@
 @testset "chainrules.jl" begin
-    @testset "BlockDiagonal" begin
-        x = [randn(1, 2), randn(2, 2)]
-        test_rrule(BlockDiagonal, x)
-    end
+    @testset for V in (Tuple, Vector)
+        @testset "BlockDiagonal" begin
+            x = V([randn(1, 2), randn(2, 2)])
+            test_rrule(BlockDiagonal, x)
+        end
 
-    @testset "Matrix" begin
-        D = BlockDiagonal([randn(1, 2), randn(2, 2)])
-        test_rrule(Matrix, D)
-    end
+        @testset "Matrix" begin
+            B = BlockDiagonal(V([randn(1, 2), randn(2, 2)]))
+            test_rrule(Matrix, B)
+        end
 
-    @testset "BlockDiagonal * Vector" begin
-        D = BlockDiagonal([rand(2, 3), rand(3, 3)])
-        v = rand(6)
-        test_rrule(*, D, v)
+        @testset "BlockDiagonal * Vector" begin
+            B = BlockDiagonal(V([rand(2, 3), rand(3, 3)]))
+            v = rand(6)
+            test_rrule(*, B, v)
+        end
     end
 end