Test weirder array types

Author: oxinabox

Project.toml

@@ -14,13 +14,15 @@ ChainRulesCore = "0.10.4"
 ChainRulesTestUtils = "0.7.9"
 Compat = "3.30"
 FiniteDifferences = "0.12.8"
+StaticArrays = "1.2"
 julia = "1"
 
 [extras]
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["ChainRulesTestUtils", "FiniteDifferences", "Random", "Test"]
+test = ["ChainRulesTestUtils", "FiniteDifferences", "Random", "StaticArrays", "Test"]

src/rulesets/Base/mapreduce.jl

@@ -27,7 +27,7 @@ function rrule(
 
     pullbacks = last.(fx_and_pullbacks)
     function sum_pullback(ȳ)
-        f̄_and_x̄s = [pullback(ȳ) for pullback in pullbacks]
+        f̄_and_x̄s = map(pullback->pullback(ȳ), pullbacks)
         # no point thunking as most of work is in f̄_and_x̄s which we need to compute for both
         f̄ = if fieldcount(typeof(f)) === 0 # Then don't need to worry about derivative wrt f
             NoTangent()

test/rulesets/Base/mapreduce.jl

@@ -30,6 +30,18 @@
 
         test_rrule(sum, abs, [-4.0 2.0; 2.0 -1.0])
         test_rrule(sum, abs, [-4.0 2.0; 2.0 -1.0]')
+
+        test_rrule(sum, abs, @SVector[1.0, -3.0])
+
+        # Make sure we preserve type for StaticArrays
+        ADviaRuleConfig = ChainRulesTestUtils.ADviaRuleConfig
+        _, pb = rrule(ADviaRuleConfig(), sum, abs, @SVector[1.0, -3.0])
+        @test pb(1.0) isa Tuple{NoTangent, NoTangent, SVector{2, Float64}}
+
+        # For structured sparse matrixes we screw it up, getting dense back
+        # see https://github.com/JuliaDiff/ChainRules.jl/issues/232 etc
+        _, pb = rrule(ADviaRuleConfig(), sum, abs, Diagonal([1.0, -3.0]))
+        @test_broken pb(1.0)[3] isa Diagonal
     end
 
     @testset "prod" begin

test/runtests.jl

@@ -10,6 +10,7 @@ using LinearAlgebra
 using LinearAlgebra.BLAS
 using LinearAlgebra: dot
 using Random
+using StaticArrays
 using Statistics
 using Test