add map/broadcast/algebra/iteration/dict interface for Grads

docs/src/utils.md

@@ -23,3 +23,32 @@ Zygote.checkpointed
 ```
 
 `Params` and `Grads` can be copied to and from arrays using the `copy!` function.
+
+## Working with Grads
+
+Map, broadcast, and iteration are supported for the dictionary-like `Grads` objects.
+These operations are value based and preserve the keys.
+
+```julia
+using Zygote, Test
+
+w, x1, x2, b = rand(2), rand(2), rand(2), rand(2)
+
+gs1 = gradient(() -> sum(tanh.(w .* x1 .+ b)), Params([w, b])) 
+gs2 = gradient(() -> sum(tanh.(w .* x2 .+ b)), Params([w, b])) 
+
+# accumulate gradients
+gs = gs1 .+ gs2
+@test gs[w] ≈ gs1[w] + gs2[w] 
+@test gs[b] ≈ gs1[b] + gs2[b] 
+
+# clip gradients
+map(x -> clamp.(x, -0.1, 0.1), gs)
+
+# clip gradients in-place
+foreach(x -> clamp!(x, -0.1, 0.1), gs)
+
+for (p, g) in pairs(gs)
+  # do something with parameter `p` and corresponding gradient `g`
+end
+```

src/compiler/interface.jl

@@ -2,6 +2,7 @@ using InteractiveUtils
 using InteractiveUtils: typesof
 using Core: Typeof
 import Base: copy!
+import Base.Broadcast: broadcasted, materialize!
 
 mutable struct Context <: AContext
   cache::Union{IdDict{Any,Any},Nothing}
@@ -139,7 +140,23 @@ end
 
 Base.show(io::IO, ps::Grads) = print(io, "Grads(...)")
 
-@forward Grads.grads Base.getindex, Base.haskey
+@forward Grads.grads  Base.haskey, Base.setindex!
+@forward Grads.params  Base.length
+
+const ADictOrGrads = Union{AbstractDict, Grads}
+
+
+# Dictionary interface.
+# Don't use the IdDict directly since it may contain some spurious pairs. 
+Base.keys(gs::Grads) = gs.params
+Base.values(gs::Grads) = (gs.grads[p] for p in gs.params)
+
+function Base.iterate(gs::Grads, state...)
+  res = iterate(gs.params, state...)
+  isnothing(res) && return nothing
+  p, next_state = res
+  return gs[p], next_state
+end
 
 function Base.getindex(gs::Grads, x)
   isbits(x) && error("Only reference types can be differentiated with `Params`.")
@@ -171,6 +188,42 @@ function copy!(x::AbstractVector,  gs::Grads)
   x
 end
 
+broadcasted(f, gs::Grads, gss::ADictOrGrads...) = map(f, gs, gss...)
+
+for op in (:*, :/)
+  @eval broadcasted(::typeof($op), a::Number, gs::Grads) = map(x -> $op(a, x), gs)
+  @eval broadcasted(::typeof($op), gs::Grads, a::Number) = map(x -> $op(x, a), gs)
+end
+
+function materialize!(gs1::Grads, gs2::Grads)
+  issetequal(gs1.params, gs2.params) || 
+    throw(ArgumentError("Expected Grads objects with the same Params."))
+  for p in gs1.params
+    gs1[p] = gs2[p]
+  end
+  return gs1
+end
+
+
+function Base.map(f, gs1::Grads, gss::ADictOrGrads...)
+  gsout = Grads(IdDict{Any,Any}(), Params(gs1.params))
+  return map!(f, gsout, gs1, gss...)
+end
+
+function Base.map!(f, gsout::Grads, gss::ADictOrGrads...)
+  all(issetequal(gsout.params, keys(gs)) for gs in gss) || 
+    throw(ArgumentError("map! expects Grads objects with the same Params."))
+  for p in gsout.params
+    gsout[p] = f((_getformap(gs, p) for gs in gss)...) 
+  end
+  return gsout
+end
+
+function _getformap(gs, p)
+  g = gs[p]
+  isnothing(g) ? fill!(similar(p), 0) : g 
+end
+
 function pullback(f, ps::Params)
   cx = Context()
   y, back = _pullback(cx, f)

test/interface.jl

@@ -1,4 +1,6 @@
-@testset "Parmas" begin
+using Zygote: Grads
+
+@testset "Params" begin
   @testset "delete!" begin
     w = rand(2,3)
     b = rand(2)
@@ -55,4 +57,84 @@
     @test ps1 == ps2 
     @test ps1 != ps3  # comparison is order dependent
   end
-end
+end
+
+@testset "Grads" begin
+  @testset "algebra" begin
+    w, b = rand(2), rand(2)
+    x1, x2 = rand(2), rand(2)
+
+    gs1 = gradient(() -> sum(w .* x1), Params([w])) 
+    gs2 = gradient(() -> sum(w .* x2), Params([w])) 
+
+    @test .- gs1 isa Grads
+    @test gs1 .- gs2 isa Grads 
+    @test .+ gs1 isa Grads
+    @test gs1 .+ gs2 isa Grads 
+    @test 2 .* gs1 isa Grads 
+    @test (2 .* gs1)[w] ≈ 2 * gs1[w]
+    @test gs1 .* 2 isa Grads 
+    @test gs1 ./ 2 isa Grads  
+    @test (gs1 .+ gs2)[w] ≈ gs1[w] .+ gs2[w] 
+
+    gs12 = gs1 .+ gs2
+    gs1 .+= gs2
+    @test gs12[w] ≈ gs1[w] 
+
+    gs3 = gradient(() -> sum(w .* x1), Params([w, b])) # grad nothing with respect to b
+    gs4 = gradient(() -> sum(w .* x2 .+ b), Params([w, b])) 
+
+    @test .- gs3 isa Grads
+    @test gs3 .- gs4 isa Grads 
+    @test .+ gs3 isa Grads
+    @test gs3 .+ gs4 isa Grads 
+    @test 2 .* gs3 isa Grads 
+    @test gs3 .* 2 isa Grads 
+    @test gs3 ./ 2 isa Grads  
+    @test (gs3 .+ gs4)[w] ≈ gs3[w] .+ gs4[w]
+    @test (gs3 .+ gs4)[b] ≈ gs4[b] 
+
+    @test gs3 .+ Dict(w => similar(w), b => similar(b)) isa Grads
+  end
+
+  @testset "map and broadcast" begin
+    w = rand(2)
+    x1 = rand(2)
+    x2 = rand(2)
+
+    gs1 = gradient(() -> sum(w .* x1), Params([w])) 
+    gs2 = gradient(() -> sum(w .* x2), Params([w])) 
+
+    @test map(x -> zeros(2), gs1) isa Grads
+
+    gs11 = map(x -> clamp.(x, -1e-5, 1e-5), gs1) 
+    @test gs11 isa Grads
+    @test all(abs.(gs11[w]) .<= 1e-5) 
+
+    @test (x -> zeros(2)).(gs1) isa Grads
+  end
+
+  @testset "dictionary interface" begin
+    w, b, x = rand(2), rand(2), rand(2)
+    ps = Params([w, b])
+    gs = gradient(() -> sum(tanh.(w .* x .+ b)), ps) 
+
+    @test issetequal(keys(gs), ps) 
+    @test length(values(gs)) == 2
+    @test length(pairs(gs)) == 2
+    k, v = first(pairs(gs))
+    @test k === first(ps) 
+    @test v === gs[first(ps)]  
+  end
+
+  @testset "iteration" begin
+    w, b, x = rand(2), rand(2), rand(2)
+    ps = Params([w, b])
+    gs = gradient(() -> sum(tanh.(w .* x .+ b)), ps) 
+
+    # value-based iteration
+    foreach(x -> clamp!(x, -1e-5, 1e-5), gs)
+    @test all(abs.(gs[w]) .<= 1e-5) 
+    @test all(abs.(gs[b]) .<= 1e-5) 
+  end
+end