add support for dicts (#122)

* add support for dicts

* aaa

* fix update

* cleanup

* more tests

* Update src/interface.jl

Co-authored-by: Brian Chen <[email protected]>

Co-authored-by: Brian Chen <[email protected]>

Author: CarloLucibello

src/backup.jl

@@ -40,7 +40,7 @@ function _setup(rule, x; cache)
       cache[x] = ℓ
     end
   else
-    map(xᵢ -> _setup(rule, xᵢ; cache), _trainable(x))
+    valuemap(xᵢ -> _setup(rule, xᵢ; cache), _trainable(x))
   end
 end
 
@@ -77,7 +77,7 @@ function _update!(tree, x; grads, params)
   haskey(params, (tree,x)) && return params[(tree,x)]
   isbits(tree) && return x  # means () is not cached, and also (((),),)
   x′, re = functor(x)
-  x′′ = re(map((tᵢ, xᵢ) -> _update!(tᵢ, xᵢ; grads, params), tree, x′))
+  x′′ = re(valuemap((tᵢ, xᵢ) -> _update!(tᵢ, xᵢ; grads, params), tree, x′))
   if ismutable(x′′)
     params[(tree,x)] = x′′
   else  # no ties to preserve between immutable structs, right?
@@ -109,7 +109,7 @@ function _grads!(dict::IdDict, tree, x, x̄s...)
   # functor(typeof(tree), base(x̄)), for things like Transpose
   x̄s′ = map(x̄ -> functor(typeof(x), base(x̄))[1], x̄s)
   x′, _ = functor(typeof(x), x)
-  foreach((tᵢ, xᵢ, x̄sᵢ...) -> _grads!(dict, tᵢ, xᵢ, x̄sᵢ...), tree, x′, x̄s′...)
+  valueforeach((tᵢ, xᵢ, x̄sᵢ...) -> _grads!(dict, tᵢ, xᵢ, x̄sᵢ...), tree, x′, x̄s′...)
 end
 
 # default all rules to first order calls
@@ -160,11 +160,22 @@ _trainable(x) = _trainable(functor(x)[1], trainable(x))
 _trainable(ch::NamedTuple, tr::NamedTuple) = merge(map(_ -> nothing, ch), tr)
 _trainable(ch::Tuple{Vararg{Any,N}}, tr::Tuple{Vararg{Any,N}}) where N = tr
 _trainable(ch::AbstractArray, tr::AbstractArray) = tr
+_trainable(ch::Dict, tr::Dict) = merge(valuemap(_ -> nothing, ch), tr)
+
 function _trainable(ch::NamedTuple, tr::Tuple)  # for old Flux-style no-names tuple
   @warn "trainable(x) should now return a NamedTuple with the field names, not a Tuple" maxlog=3
   map(c -> c in tr ? c : nothing, ch)
 end
 
+
+valuemap(f, x...) = map(f, x...)
+valuemap(f, x::Dict, ys...) = Dict(k => f(v, (get(y, k, nothing) for y in ys)...) for (k,v) in x)
+valueforeach(f, x...) = foreach(f, x...)
+valueforeach(f, x::Dict, ys...) = foreach(pairs(x)) do (k, v)
+  f(v, (get(y, k, nothing) for y in ys)...)
+end
+
+
 ###
 ### rule definition helpers
 ###

src/interface.jl

@@ -99,6 +99,36 @@ y2z(x) = x
       @test isnan(m3n.γ[3])
     end
 
+    @testset "Dict support" begin
+      @testset "simple dict" begin
+        d = Dict(:a => [1.0,2.0], :b => [3.0,4.0], :c => 1)
+        s = Optimisers.setup(AdamW(0.1), d)
+        @test s isa Dict{Symbol, <:Any}
+        @test s[:a] isa Optimisers.Leaf
+        @test s[:b] isa Optimisers.Leaf
+        @test s[:c] === ()
+        loss(model) = sum(abs2, model[:a])
+        g = gradient(loss, d)[1]
+        s2, d2 = Optimisers.update(s, d, g)
+        @test s2 isa Dict{Symbol, <:Any}
+        @test d2 isa Dict{Symbol, <:Any}
+        @test d2[:a] == [0.9, 1.9]
+        @test d2[:b] == [3, 4]
+        @test d2[:c] == 1
+      end
+
+      @testset "nested dict" begin
+        d = Dict(1 => [1.0,2.0], 2 => Dict("a" => (; c=[3.0,4.0]), "b" => 1))
+        s = Optimisers.setup(AdamW(0.1), d)
+        @test s[2]["a"].c isa Optimisers.Leaf
+        g = gradient(d -> sum(d[2]["a"].c), d)[1]
+        s2, d2 = Optimisers.update(s, d, g)
+        @test d2[2]["a"].c == [2.9, 3.9]
+        @test d2[1] == [1, 2]
+        @test d2[2]["b"] == 1
+      end
+    end
+
     @testset "OptimiserChain" begin
       x = [1, 10, 100.0]; dx = [1, 2, 3.0];
       @test Optimisers.update(Optimisers.setup(WeightDecay(0.1), x), x, dx)[2] ≈ [1-0.1-1, 10-1-2, 100-10-3]