simple rule for mapfoldl

Author: mcabbott

Project.toml

@@ -14,7 +14,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
-ChainRulesCore = "1.12"
+ChainRulesCore = "1.15.3"
 ChainRulesTestUtils = "1.5"
 Compat = "3.42.0, 4"
 FiniteDifferences = "0.12.20"

src/rulesets/Base/mapreduce.jl

@@ -417,6 +417,25 @@ end
 end
 
 #####
+##### `mapfoldl(f, g, ::Tuple)`
+#####
+
+# For tuples there should be no harm in handling `map` first.
+# This will also catch `mapreduce`.
+
+function rrule(
+        cfg::RuleConfig{>:HasReverseMode}, ::typeof(Base.mapfoldl_impl), f::F, op::G, init, x::Tuple;
+    ) where {F,G}
+    y, backmap = rrule(cfg, map, f, x)
+    z, backred = rrule(cfg, Base.mapfoldl_impl, identity, op, init, y)
+    function mapfoldl_pullback_tuple(dz)
+        _, _, dop, dinit, dy = backred(dz)
+        _, df, dx = backmap(dy)
+        return (NoTangent(), df, dop, dinit, dx)
+    end
+    return z, mapfoldl_pullback_tuple
+end
+
 #####
 ##### `foldl(f, ::Tuple)`
 #####

test/rulesets/Base/mapreduce.jl

@@ -303,6 +303,11 @@ const _INIT = Base._InitialValue()
         test_rrule(mapfoldl_impl, identity, /, _INIT, Tuple(1 .+ rand(5)))
         test_rrule(mapfoldl_impl, identity, *, 1+rand(), Tuple(rand(ComplexF64, 5)))
     end
+    @testset "mapfoldl(f, g, ::Tuple)" begin
+        test_rrule(mapfoldl_impl, cbrt, /, _INIT, Tuple(1 .+ rand(5)), check_inferred=false)
+        test_rrule(mapfoldl_impl, abs2, *, 1+rand(), Tuple(rand(ComplexF64, 5)), check_inferred=false)
+        # TODO make the `map(f, ::Tuple)` rule infer better!
+    end
 end
 
 @testset "Accumulations" begin