tweak, simplify recursion

Author: mcabbott

src/adjust.jl

@@ -31,16 +31,16 @@ julia> s.x
 (Leaf(Momentum{Float32}(0.01, 0.9), [0.0]), ())
 ```
 """
-freeze!(tree) = (fmapstructure(freeze!, tree; exclude = x -> x isa Leaf); nothing)
+freeze!(tree) = foreach(freeze!, tree)
 freeze!(ℓ::Leaf) = (ℓ.frozen = true; nothing)
 
 """
     Optimisers.thaw!(tree)
 
 Un-does [`freeze!`](@ref Optimisers.freeze!) for all parameters,
-mutating every `Leaf(rule, state, true)` to `Leaf(rule, state, false)`.
+mutating every `Leaf(rule, state, frozen = true)` to `Leaf(rule, state, frozen = false)`.
 """
-thaw!(tree) = (fmapstructure(thaw!, tree; exclude = x -> x isa Leaf); nothing)
+thaw!(tree) = foreach(thaw!, tree)
 thaw!(ℓ::Leaf) = (ℓ.frozen = false; nothing)
 
 freeze!(::Union{Number, AbstractArray{<:Number}}) = throw(ArgumentError(

src/interface.jl

@@ -49,7 +49,7 @@ function Base.show(io::IO, ℓ::Leaf; colour = ℓ.frozen ? :cyan : :green)
   str = sprint(show, ℓ.rule; context = ioc)
   printstyled(io, "Leaf(", str, ", "; color = colour)
   show(ioc, ℓ.state)
-  printstyled(io, ℓ.frozen ? ", frozen=true)" : ")"; color = colour)
+  printstyled(io, ℓ.frozen ? ", frozen = true)" : ")"; color = colour)
 end
 
 ###