Merge #924

924: Throw error when trying to accum two NamedTuples with different keys r=DhairyaLGandhi a=mzgubic

Would have helped here:
#922 (comment)



Co-authored-by: Miha Zgubic <[email protected]>

Project.toml

@@ -1,6 +1,6 @@
 name = "Zygote"
 uuid = "e88e6eb3-aa80-5325-afca-941959d7151f"
-version = "0.6.4"
+version = "0.6.5"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/lib/lib.jl

@@ -16,7 +16,10 @@ accum(x::Tuple, y::Tuple) = accum.(x, y)
 accum(x::AbstractArray, y::AbstractArray) = accum.(x, y)
 
 @generated function accum(x::NamedTuple, y::NamedTuple)
-  grad(x) = x in fieldnames(y) ? :(y.$x) : :nothing
+  # assumes that y has no keys apart from those also in x
+  fieldnames(y) ⊆ fieldnames(x) || throw(ArgumentError("$y keys must be a subset of $x keys"))
+
+  grad(field) = field in fieldnames(y) ? :(y.$field) : :nothing
   Expr(:tuple, [:($f=accum(x.$f, $(grad(f)))) for f in fieldnames(x)]...)
 end
 

test/lib/lib.jl

@@ -0,0 +1,8 @@
+@testset "lib.jl" begin
+    @testset "accum" begin
+        t1 = (a=1, b=2, c=3)
+        t2 = (a=1, b=2)
+        @test Zygote.accum(t1, t2) == (a = 2, b = 4, c = 3)
+        @test_throws ArgumentError Zygote.accum(t2, t1)
+    end
+end

test/runtests.jl

@@ -22,8 +22,9 @@ end
   include("utils.jl")
 end
 
-@testset "lib/number" begin
+@testset "lib" begin
   include("lib/number.jl")
+  include("lib/lib.jl")
 end
 
 @testset "Features" begin