Rewrite as a generated functor

Author: oxinabox

src/rulesets/Base/base.jl

@@ -295,17 +295,21 @@ end
 ####
 #### merge
 ####
-
-function rrule(::typeof(merge), nt1::NamedTuple{F1}, nt2::NamedTuple{F2}) where {F1,F2}
-    y = merge(nt1, nt2)
-    function merge_pullback(dy)
-        dnt1 = Tangent{typeof(nt1)}(;
-            (f1 => (f1 in F2 ? ZeroTangent() : getproperty(dy, f1)) for f1 in F1)...
-        )
-        dnt2 = Tangent{typeof(nt2)}(; (f2 => getproperty(dy, f2) for f2 in F2)...)
+# need to work around inability to return closures from generated functions
+struct MergePullback{T1, T2}
+end
+(this::MergePullback)(dy::AbstractThunk) = this(unthunk(dy))
+(::MergePullback)(x::AbstractZero) = (NoTangent(), x, x)
+@generated function(::MergePullback{T1,T2})(dy::Tangent) where {F1,T1<:NamedTuple{F1},F2,T2<:NamedTuple{F2}}
+    _getproperty_kwexpr(key) = :($key = getproperty(dy, $(Meta.quot(key))))
+    quote
+        dnt1 = Tangent{T1}(; $(map(_getproperty_kwexpr, setdiff(F1, F2))...))
+        dnt2 = Tangent{T2}(; $(map(_getproperty_kwexpr, F2)...))
         return (NoTangent(), dnt1, dnt2)
     end
-    merge_pullback(dy::AbstractThunk) = merge_pullback(unthunk(dy))
-    merge_pullback(x::AbstractZero) = (NoTangent(), x, x)
-    return y, merge_pullback
+end
+
+function rrule(::typeof(merge), nt1::T1, nt2::T2) where {T1<:NamedTuple, T2<:NamedTuple}
+    y = merge(nt1, nt2)
+    return y, MergePullback{T1,T2}()
 end

test/rulesets/Base/base.jl

@@ -259,7 +259,7 @@ end
     end
 
     @testset "merge NamedTuple" begin
-        test_rrule(merge, (; a=1.0), (; b=2.0); check_inferred=false)
-        test_rrule(merge, (; a=1.0), (; a=2.0); check_inferred=false)
+        test_rrule(merge, (;a=1.0), (;b=2.0))
+        test_rrule(merge, (;a=1.0), (;a=2.0))
     end
 end