Merge pull request #795 from JuliaDiff/ox/merge

Add rrule for NamedTuple merge

Author: oxinabox

src/rulesets/Base/base.jl

@@ -291,3 +291,26 @@ function rrule(config::RuleConfig{>:HasReverseMode}, ::typeof(task_local_storage
     end
     return y, task_local_storage_pullback
 end
+
+####
+#### merge
+####
+# 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
+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

@@ -258,4 +258,9 @@ end
             test_rrule(map, Multiplier(4.5), (6.7, 8.9), (0.1, 0.2, 0.3), check_inferred=false)
         end
     end
+
+    @testset "merge NamedTuple" begin
+        test_rrule(merge, (; a=1.0), (; b=2.0))
+        test_rrule(merge, (; a=1.0), (; a=2.0))
+    end
 end