decouple from adjust, tweak words

Author: mcabbott

docs/src/index.md

@@ -142,7 +142,7 @@ st = Optimisers.setup(DecayDescent(0.1), Layer(3))
 
 To temporarily prevent training from affecting some parameters,
 use [freeze!](@ref Optimisers.freeze!) and `thaw!`.
-They work by mutating all `Leaf`s of the state tree, within the given limb:
+They work by mutating all `Leaf`s of the state tree, or part of it.
 
 ```julia
 using Flux, Optimisers

src/adjust.jl

@@ -1,5 +1,5 @@
 ###
-### freeze!
+### freezing
 ###
 
 """
@@ -12,8 +12,6 @@ Can be applied to the state corresponding to only part of a model,
 for instance with `model::Chain`, to freeze `model.layers[1]` you
 should call `freeze!(tree.layers[1])`.
 
-Also prevents [`adjust`](@ref Optimisers.adjust) from changing the rule's parameters.
-
 # Example
 ```jldoctest
 julia> m = (x = ([1.0], 2.0), y = [3.0]);
@@ -66,8 +64,6 @@ through training.
 
 To change just the learning rate, provide a number `η::Real`.
 
-Does not affect any frozen parameters, set by [`freeze!`](@ref Optimisers.freeze!).
-
 # Example
 ```jldoctest
 julia> m = (vec = rand(Float32, 2), fun = sin);
@@ -107,8 +103,8 @@ adjust(tree; kw...) = map(st -> adjust(st; kw...), tree)
 adjust(::Nothing, ::Real) = nothing
 adjust(::Nothing; kw...) = nothing
 
-adjust(ℓ::Leaf, eta::Real) = ℓ.frozen ? ℓ : Leaf(adjust(ℓ.rule, eta), ℓ.state)
-adjust(ℓ::Leaf; kw...) = ℓ.frozen ? ℓ : Leaf(adjust(ℓ.rule; kw...), ℓ.state)
+adjust(ℓ::Leaf, eta::Real) = Leaf(adjust(ℓ.rule, eta), ℓ.state, ℓ.frozen)
+adjust(ℓ::Leaf; kw...) = Leaf(adjust(ℓ.rule; kw...), ℓ.state, ℓ.frozen)
 
 
 """

test/runtests.jl

@@ -232,7 +232,7 @@ end
       st = Optimisers.adjust(st, 0.2)
       Optimisers.thaw!(st)
       st, m = Optimisers.update(st, m, (x=[1,10], y=([100,1000], nothing)));
-      @test m.y[1] ≈ [-7.0, -96.0]
+      @test m.y[1] ≈ [-17.0, -196.0]
       @test m.x ≈ [0.7, -1.0]
 
       @test_throws ArgumentError Optimisers.freeze!(m)