Allow types from ChainRules (#41)

* using ChainRulesCore

* version

* indents

* add some tests

* no weird indents

Author: mcabbott

Project.toml

@@ -4,20 +4,21 @@ authors = ["Mike J Innes <mike.j.innes@gmail.com>"]
 version = "0.2.0"
 
 [deps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
+ChainRulesCore = "1"
 Functors = "0.2.7"
 julia = "1.6"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "ChainRulesCore", "StaticArrays", "Zygote"]
+test = ["Test", "StaticArrays", "Zygote"]

src/interface.jl

@@ -1,4 +1,9 @@
 
+using ChainRulesCore: canonicalize, backing, Tangent, AbstractZero
+base(dx::Tangent) = backing(canonicalize(dx))
+base(dx) = dx
+const Zero = Union{Nothing, AbstractZero}  # Union{Zygote, Diffractor}
+
 struct Leaf{R,S}
   rule::R
   state::S
@@ -18,26 +23,21 @@ end
 
 subtract!(x, x̄) = iswriteable(x) ? (x .= x .- x̄) : (x .- x̄)
 
+update!(::Nothing, x, ::Zero...) = nothing, x
 update!(::Nothing, x, x̄s...) = nothing, x
 
+update!(ℓ::Leaf, x, ::Zero...) = ℓ, x
 function update!(ℓ::Leaf, x, x̄s...)
-  if all(isnothing, x̄s)
-    return ℓ, x
-  else
-    s′, x̄′ = apply!(ℓ.rule, ℓ.state, x, x̄s...)
-    return Leaf(ℓ.rule, s′), subtract!(x, x̄′)
-  end
+  s′, x̄′ = apply!(ℓ.rule, ℓ.state, x, base.(x̄s)...)
+  Leaf(ℓ.rule, s′), subtract!(x, x̄′)
 end
 
+update!(tree, x, ::Zero...) = tree, x
 function update!(tree, x, x̄s...)
-  if all(isnothing, x̄s)
-    return tree, x
-  else
-    x̄s′ = map(x̄ -> functor(typeof(x), x̄)[1], x̄s)
-    x′, re = functor(typeof(x), x)
-    xtree = map((stᵢ, xᵢ, x̄sᵢ...) -> update!(stᵢ, xᵢ, x̄sᵢ...), tree, x′, x̄s′...)
-    return map(first, xtree), re(map(last, xtree))
-  end
+  x̄s′ = map(x̄ -> functor(typeof(x), base(x̄))[1], x̄s)
+  x′, re = functor(typeof(x), x)
+  xtree = map((stᵢ, xᵢ, x̄sᵢ...) -> update!(stᵢ, xᵢ, x̄sᵢ...), tree, x′, x̄s′...)
+  map(first, xtree), re(map(last, xtree))
 end
 
 function update(tree, x, x̄s...)

test/runtests.jl

@@ -31,6 +31,10 @@ Optimisers.trainable(x::TwoThirds) = (a = x.a,)
       s3, m3 = Optimisers.update!(s, m, g)
       @test objectid(m3[1]) == mid
       @test m3[1] ≈ [1,2] .- 0.1 .* [25, 33]
+
+      g4 = Tangent{typeof(m)}(g...)
+      s4, m4 = Optimisers.update!(s, ([1.0, 2.0],), g4)
+      @test m4[1] ≈ [1,2] .- 0.1 .* [25, 33]
     end
 
     @testset "gradient clipping" begin
@@ -74,6 +78,11 @@ Optimisers.trainable(x::TwoThirds) = (a = x.a,)
       @test mf2.x == [1,2]
       @test mf2.y == (a = sin, b = [2.9, 3.9], c = 5)
 
+      gf3 = Tangent{typeof(mf)}(; x = NoTangent(), y = Tangent{typeof(mf.y)}(; a = NoTangent(), b = [1,1], c = 1))
+      _, mf3 = Optimisers.update(sf, mf, gf3)  # the same, but with ChainRules types
+      @test mf3.x == [1,2]
+      @test mf3.y == (a = sin, b = [2.9, 3.9], c = 5)
+
       # TwoThirds has functor a,c only, and trainable a only
       mt = TwoThirds(Float32[1,2], Float32[3,4], Float32[5,6])
       mt10 = fmap(x -> 10x, mt)