accumulate NamedTuple + Tangent

Author: mcabbott

src/runtime.jl

@@ -13,3 +13,17 @@ end
 @Base.constprop :aggressive accum(a::NoTangent, b) = b
 @Base.constprop :aggressive accum(a, b::NoTangent) = a
 @Base.constprop :aggressive accum(a::NoTangent, b::NoTangent) = NoTangent()
+
+using ChainRulesCore: Tangent, backing
+
+function accum(x::Tangent{T}, y::NamedTuple) where T
+  # @warn "gradient is both a Tangent and a NamedTuple" x y
+  z = accum(backing(x), y)
+  Tangent{T,typeof(z)}(z)
+end
+accum(x::NamedTuple, y::Tangent) = accum(y, x)
+
+function accum(x::Tangent{T}, y::Tangent) where T
+  z = accum(backing(x), backing(y))
+  Tangent{T,typeof(z)}(z)
+end

test/runtests.jl

@@ -162,15 +162,7 @@ end
 # Make sure that there's no infinite recursion in kwarg calls
 g_kw(;x=1.0) = sin(x)
 f_kw(x) = g_kw(;x)
-@test bwd(f_kw)(1.0) == bwd(sin)(1.0)  broken=true
-#=
-MethodError: no method matching +(::Tangent{var"#g_kw#47"{var"#g_kw#11#48"}, NamedTuple{(Symbol("#g_kw#11"),), Tuple{ZeroTangent}}}, ::Tangent{Diffractor.KwFunc{var"#g_kw#47"{var"#g_kw#11#48"}, var"#g_kw#47##kw"}, NamedTuple{(:kwf,), Tuple{ZeroTangent}}})
-...
-  [2] elementwise_add(a::NamedTuple{(:contents,), Tuple{Tangent{var"#g_kw#47"{var"#g_kw#11#48"}, NamedTuple{(Symbol("#g_kw#11"),), Tuple{ZeroTangent}}}}}, b::NamedTuple{(:contents,), Tuple{Tangent{Diffractor.KwFunc{var"#g_kw#47"{var"#g_kw#11#48"}, var"#g_kw#47##kw"}, NamedTuple{(:kwf,), Tuple{ZeroTangent}}}}})
-    @ ChainRulesCore ~/.julia/packages/ChainRulesCore/ctmSK/src/tangent_types/tangent.jl:287
-  [3] +(a::Tangent{Core.Box, NamedTuple{(:contents,), Tuple{Tangent{var"#g_kw#47"{var"#g_kw#11#48"}, NamedTuple{(Symbol("#g_kw#11"),), Tuple{ZeroTangent}}}}}}, b::Tangent{Core.Box, NamedTuple{(:contents,), Tuple{Tangent{Diffractor.KwFunc{var"#g_kw#47"{var"#g_kw#11#48"}, var"#g_kw#47##kw"}, NamedTuple{(:kwf,), Tuple{ZeroTangent}}}}}})
-    @ ChainRulesCore ~/.julia/packages/ChainRulesCore/ctmSK/src/tangent_arithmetic.jl:130
-=#
+@test bwd(f_kw)(1.0) == bwd(sin)(1.0)
 
 function f_crit_edge(a, b, c, x)
     # A function with two critical edges. This used to trigger an issue where