Merge pull request #992 from DhairyaLGandhi/dg/941

Differentiate `push!` with implicit Params

Author: CarloLucibello

src/compiler/interface.jl

@@ -65,11 +65,15 @@ Base.adjoint(f::Function) = x -> gradient(f, x)[1]
 
 # TODO store ids only
 struct Params
-  order::Buffer{Any, Vector{Any}}
+  order::Buffer # {Any, Vector{Any}}
   params::IdSet{Any}
-  Params() = new(Buffer([], false), IdSet())
 end
 
+Params() = Params(Buffer([], false), IdSet())
+Params(xs) = Params(Buffer(xs, false), IdSet(xs))
+Params(ps::Params) = ps
+Params(xs::Tuple) = Params(collect(xs))
+
 @forward Params.order Base.iterate, Base.length, Base.getindex
 @forward Params.params Base.in
 
@@ -103,6 +107,20 @@ function Base.push!(ps::Params, x)
   return ps
 end
 
+@adjoint! function Base.push!(xs::IdSet, x...)
+  l = length(x)
+  push!(xs, x...), Δ -> begin
+    (Δ, ntuple(_ -> nothing, l)...)
+  end
+end
+
+@adjoint! function Base.push!(xs::Params, x::AbstractArray{T}...) where T
+  sz_x = size.(x)
+  push!(xs, x...), Δ -> begin
+    (Δ, map(x -> Ones{T}(x...), sz_x)...)
+  end
+end
+
 Base.push!(ps::Params, x...) = (foreach(x -> push!(ps, x), x); ps)
 
 function Base.delete!(ps::Params, x)
@@ -114,8 +132,6 @@ function Base.delete!(ps::Params, x)
   return ps
 end
 
-Params(xs) = push!(Params(), xs...)
-
 Base.Broadcast.broadcasted(f, ps::Params) = broadcasted(f, ps.order)
 
 Base.:(==)(x::Params, y::Params) = x.order.data == y.order.data

test/interface.jl

@@ -163,4 +163,44 @@ end
     @test all(abs.(gs[w]) .<= 1e-5) 
     @test all(abs.(gs[b]) .<= 1e-5) 
   end
+
+  @testset "Params nesting" begin
+    struct Dense{F,T,S}
+      W::T
+      b::S
+      σ::F
+    end
+  
+    (d::Dense)(x) = d.σ.(d.W * x .+ d.b)
+    d = Dense(ones(Float32, 3,3), zeros(Float32, 3), identity)
+    ps = Zygote.Params([d.W, d.b])
+    r = ones(Float32, 3,3)
+  
+    gs = gradient(ps) do
+      p, pb = pullback(ps) do
+        sum(d(r))
+      end
+      g = pb(p)
+      sum(g[d.W]) # + sum(g[d.b])
+    end
+  
+    @test gs[d.W] ≈ fill(81f0, (3,3))
+  
+    # Test L2
+    l2g = gradient(ps) do
+      sum(sum(x .^ 2) for x in ps)
+    end
+    @test l2g[d.W] ≈ fill(2.f0, size(d.W))
+    @test l2g[d.b] ≈ fill(0.f0, size(d.b))
+
+    # Can be safely removed - creating Params within
+    # gradient calls may break between releases.
+    sgs = gradient(ps) do
+      sum(sum(x) for x in Zygote.Params([d.W, d.b]))
+    end
+    @test sgs[d.W] ≈ fill(1.f0, size(d.W))
+    @test sgs[d.b] ≈ fill(1.f0, size(d.b))
+  end
+
+
 end