make a keyword frozen, and print it

Author: mcabbott

src/adjust.jl

@@ -22,13 +22,13 @@ julia> Optimisers.update!(s, m, (x = ([pi], 10pi), y = [100pi]));  # with fake g
 julia> m
 (x = ([1.0], 2.0), y = [-0.14159258336972558])
 
-julia> s  # Leaf(..., true) means frozen
-(x = (Leaf(Momentum{Float32}(0.01, 0.9), [0.0], true), ()), y = Leaf(Momentum{Float32}(0.01, 0.9), [3.14159]))
+julia> s
+(x = (Leaf(Momentum{Float32}(0.01, 0.9), [0.0], frozen=true), ()), y = Leaf(Momentum{Float32}(0.01, 0.9), [3.14159]))
 
 julia> Optimisers.thaw!(s)
 
-julia> s.x[1]
-Leaf(Momentum{Float32}(0.01, 0.9), [0.0])
+julia> s.x
+(Leaf(Momentum{Float32}(0.01, 0.9), [0.0]), ())
 ```
 """
 freeze!(tree) = (fmapstructure(freeze!, tree; exclude = x -> x isa Leaf); nothing)

src/interface.jl

@@ -15,7 +15,7 @@ mutable struct Leaf{R,S}  # mutable so that its identity encodes parameter shari
   state::S
   frozen::Bool  # ... and to allow freeze! to act on this.
 end
-Leaf(rule, state) = Leaf(rule, state, false)
+Leaf(rule, state; frozen::Bool = false) = Leaf(rule, state, frozen)
 
 @functor Leaf
 
@@ -44,12 +44,12 @@ function _setup(rule, x; cache)
   end
 end
 
-function Base.show(io::IO, ℓ::Leaf)  # show method is mostly to hide its long type!
+function Base.show(io::IO, ℓ::Leaf; colour = ℓ.frozen ? :cyan : :green)
   ioc = IOContext(io, :compact => true)
-  print(ioc, "Leaf(", ℓ.rule, ", ")
+  str = sprint(show, ℓ.rule; context = ioc)
+  printstyled(io, "Leaf(", str, ", "; color = colour)
   show(ioc, ℓ.state)
-  ℓ.frozen && print(ioc, ", true")
-  print(ioc, ")")
+  printstyled(io, ℓ.frozen ? ", frozen=true)" : ")"; color = colour)
 end
 
 ###