ensure that update without a gradient is an error (#50)

Author: mcabbott

src/interface.jl

@@ -23,16 +23,16 @@ end
 
 subtract!(x, x̄) = iswriteable(x) ? (x .= x .- x̄) : (x .- x̄)
 
-update!(::Nothing, x, ::Zero...) = nothing, x
+update!(::Nothing, x, ::Zero, ::Zero...) = nothing, x
 update!(::Nothing, x, x̄s...) = nothing, x
 
-update!(ℓ::Leaf, x, ::Zero...) = ℓ, x
+update!(ℓ::Leaf, x, ::Zero, ::Zero...) = ℓ, x
 function update!(ℓ::Leaf, x, x̄s...)
   s′, x̄′ = apply!(ℓ.rule, ℓ.state, x, base.(x̄s)...)
   Leaf(ℓ.rule, s′), subtract!(x, x̄′)
 end
 
-update!(tree, x, ::Zero...) = tree, x
+update!(tree, x, ::Zero, ::Zero...) = tree, x
 function update!(tree, x, x̄s...)
   x̄s′ = map(x̄ -> functor(typeof(x), base(x̄))[1], x̄s)
   x′, re = functor(typeof(x), x)

test/runtests.jl

@@ -116,6 +116,16 @@ Optimisers.trainable(x::TwoThirds) = (a = x.a,)
       @test Optimisers.update!(s, m, g...)[2] isa Foo
     end
 
+    @testset "forgotten gradient" begin
+      x = [1.0, 2.0]
+      sx = Optimisers.setup(Descent(), x)
+      @test_throws MethodError Optimisers.update(sx, x)
+
+      m = (x = x, y = sin)
+      sm = Optimisers.setup(Descent(), m)
+      @test_throws MethodError Optimisers.update(sm, m)
+    end
+
     @testset "broadcasting macros" begin
       x = [1.0, 2.0]; y = [3,4]; z = [5,6]
       @test (@lazy x + y * z) isa Broadcast.Broadcasted