Dictionary differentials (#183)

* Dictionary differentials

* Removes unnecessary space

* Tweaks comment

* Update docs Manifest

* Bump patch

Author: willtebbutt

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "0.9.4"
+version = "0.9.5"
 
 [deps]
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"

docs/Manifest.toml

@@ -7,7 +7,7 @@ uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 deps = ["MuladdMacro"]
 path = ".."
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "0.7.0"
+version = "0.9.3"
 
 [[Dates]]
 deps = ["Printf"]
@@ -40,6 +40,7 @@ uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 version = "0.21.0"
 
 [[LibGit2]]
+deps = ["Printf"]
 uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"
 
 [[Libdl]]
@@ -67,7 +68,7 @@ uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
 version = "0.3.12"
 
 [[Pkg]]
-deps = ["Dates", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Test", "UUIDs"]
+deps = ["Dates", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "UUIDs"]
 uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 
 [[Printf]]

src/differential_arithmetic.jl

@@ -115,6 +115,7 @@ function Base.:+(a::P, d::Composite{P}) where P
         return construct(P, net_backing)
     end
 end
+Base.:+(a::Dict, d::Composite{P}) where {P} = merge(+, a, backing(d))
 Base.:+(a::Composite{P}, b::P) where P = b + a
 
 # We intentionally do not define, `Base.*(::Composite, ::Composite)` as that is not meaningful

src/differentials/composite.jl

@@ -22,7 +22,7 @@ To make a `Composite` have all the fields of the primal the [`canonicalize`](@re
 function is provided.
 """
 struct Composite{P, T} <: AbstractDifferential
-    # Note: If T is a Tuple, then P is also a Tuple
+    # Note: If T is a Tuple/Dict, then P is also a Tuple/Dict
     # (but potentially a different one, as it doesn't contain differentials)
     backing::T
 end
@@ -41,6 +41,12 @@ function Composite{P}() where P<:Tuple
     return Composite{P, typeof(backing)}(backing)
 end
 
+function Composite{P}(d::Dict) where {P<:Dict}
+    return Composite{P, typeof(d)}(d)
+end
+
+Base.:(==)(a::Composite, b::Composite) = backing(a) == backing(b)
+
 function Base.show(io::IO, comp::Composite{P}) where P
     print(io, "Composite{")
     show(io, P)
@@ -51,6 +57,7 @@ end
 
 Base.convert(::Type{<:NamedTuple}, comp::Composite{<:Any, <:NamedTuple}) = backing(comp)
 Base.convert(::Type{<:Tuple}, comp::Composite{<:Any, <:Tuple}) = backing(comp)
+Base.convert(::Type{<:Dict}, comp::Composite{<:Dict, <:Dict}) = backing(comp)
 
 Base.getindex(comp::Composite, idx) = getindex(backing(comp), idx)
 Base.getproperty(comp::Composite, idx::Int) = getproperty(backing(comp), idx)  # for Tuple
@@ -79,6 +86,9 @@ function Base.map(f, comp::Composite{P, <:NamedTuple{L}}) where{P, L}
     named_vals = NamedTuple{L, typeof(vals)}(vals)
     return Composite{P, typeof(named_vals)}(named_vals)
 end
+function Base.map(f, comp::Composite{P, <:Dict}) where {P<:Dict}
+    return Composite{P}(Dict(k => f(v) for (k, v) in backing(comp)))
+end
 
 Base.conj(comp::Composite) = map(conj, comp)
 
@@ -97,6 +107,7 @@ primal types.
 """
 backing(x::Tuple) = x
 backing(x::NamedTuple) = x
+backing(x::Dict) = x
 backing(x::Composite) = getfield(x, :backing)
 
 function backing(x::T)::NamedTuple where T
@@ -235,6 +246,7 @@ function elementwise_add(a::NamedTuple{an}, b::NamedTuple{bn}) where {an, bn}
     end
 end
 
+elementwise_add(a::Dict, b::Dict) = merge(+, a, b)
 
 struct PrimalAdditionFailedException{P} <: Exception
     primal::P

test/differentials/composite.jl

@@ -26,6 +26,7 @@ end
     @testset "convert" begin
         @test convert(NamedTuple, Composite{Foo}(x=2.5)) == (; x=2.5)
         @test convert(Tuple, Composite{Tuple{Float64,}}(2.0)) == (2.0,)
+        @test convert(Dict, Composite{Dict}(Dict(4 => 3))) == Dict(4 => 3)
     end
 
     @testset "indexing, iterating, and properties" begin
@@ -58,15 +59,21 @@ end
             conj(Composite{Tuple{Float64,}}(2.0+3.0im)),
             Composite{Tuple{Float64,}}(2.0-3.0im)
         )
+        @test ==(
+            conj(Composite{Dict}(Dict(4 => 2.0 + 3.0im))),
+            Composite{Dict}(Dict(4 => 2.0 + -3.0im)),
+        )
     end
 
     @testset "extern" begin
         @test extern(Composite{Foo}(x=2.0)) == (;x=2.0)
         @test extern(Composite{Tuple{Float64,}}(2.0)) == (2.0,)
+        @test extern(Composite{Dict}(Dict(4 => 3))) == Dict(4 => 3)
 
         # with differentials on the inside
         @test extern(Composite{Foo}(x=@thunk(0+2.0))) == (;x=2.0)
         @test extern(Composite{Tuple{Float64,}}(@thunk(0+2.0))) == (2.0,)
+        @test extern(Composite{Dict}(Dict(4 => @thunk(3)))) == Dict(4 => 3)
     end
 
     @testset "canonicalize" begin
@@ -110,6 +117,13 @@ end
                 Composite{typeof(nt2)}(; nt2...)
             ) == Composite{typeof(nt_sum)}(; nt_sum...)
         end
+
+        @testset "Dicts" begin
+            d1 = Composite{Dict}(Dict(4 => 3.0, 3 => 2.0))
+            d2 = Composite{Dict}(Dict(4 => 3.0, 2 => 2.0))
+            d_sum = Composite{Dict}(Dict(4 => 3.0 + 3.0, 3 => 2.0, 2 => 2.0))
+            @test d1 + d2 == d_sum
+        end
     end
 
     @testset "+ with Primals" begin
@@ -133,6 +147,11 @@ end
             @test Composite{typeof(nty)}(; nty...) + nty == (; a=3.0, b=1.0)
         end
 
+        @testset "Dicts" begin
+            d_primal = Dict(4 => 3.0, 3 => 2.0)
+            d_tangent = Composite{typeof(d_primal)}(Dict(4 =>5.0))
+            @test d_primal + d_tangent == Dict(4 => 3.0 + 5.0, 3 => 2.0)
+        end
     end
 
     @testset "+ with Primals, with inner constructor" begin
@@ -189,6 +208,9 @@ end
             == Composite{Tuple{Float64, Float64}}(4.0, 8.0)
             == Composite{Tuple{Float64, Float64}}(2.0, 4.0) * 2
         )
+        d = Composite{Dict}(Dict(4 => 3.0))
+        two_d = Composite{Dict}(Dict(4 => 2 * 3.0))
+        @test 2 * d == two_d == d * 2
     end
 
     @testset "show" begin