Tweak printing of composites

Author: adrhill

src/lrp/show.jl

@@ -11,38 +11,52 @@ _range_string(::LastLayerTypeRule) = "last layer"
 _range_string(r::FirstNTypeRule) = "layers $(1:r.n)"
 _range_string(r::LastNTypeRule) = "last $(r.n) layers"
 
+const COLOR_COMMENT = :light_black
+const COLOR_RULE = :yellow
+const COLOR_TYPE = :blue
+const COLOR_RANGE = :green
+
+typename(x) = string(nameof(typeof(x)))
+
 function Base.show(io::IO, m::MIME"text/plain", c::Composite)
-    println(io, "Composite(")
+    println(io, "Composite", "(")
     for p in c.primitives
         _show_primitive(io, p, 2)
     end
-    return println(io, ")")
+    println(io, ")")
+    return nothing
 end
 
 function _show_primitive(io::IO, r::AbstractRulePrimitive, indent::Int=0)
-    print(io, " "^indent, nameof(typeof(r)), ": ")
-    printstyled(io, _range_string(r); color=:blue)
+    print(io, " "^indent, typename(r), ": ")
+    printstyled(io, _range_string(r); color=COLOR_RANGE)
     print(io, " => ")
-    printstyled(io, r.rule; color=:yellow)
+    printstyled(io, r.rule; color=COLOR_RULE)
     println(io, ",")
     return nothing
 end
 
-function _show_primitive(io::IO, r::AbstractCompositePrimitive, indent::Int=0)
-    print(io, " "^indent, nameof(typeof(r)), "(")
-    printstyled(io, "  # on ", _range_string(r); color=:light_black)
+function _show_primitive(io::IO, r::AbstractTypeRulePrimitive, indent::Int=0)
+    print(io, " "^indent, rpad(typename(r) * "(", 20))
+    printstyled(io, "# on ", _range_string(r); color=COLOR_COMMENT)
     println(io)
     _print_rules(io, r.map, indent)
+    println(io, " "^(indent), "),")
     return nothing
 end
 
 function _print_rules(io::IO, rs, indent::Int=0)
     for r in rs
-        printstyled(io, " "^(indent + 2), r.first; color=:blue)
-        print(io, " => ")
-        printstyled(io, r.second; color=:yellow)
+        print(io, " "^(indent + 2))
+        _print_type_rule(io, r)
         println(io, ",")
     end
-    println(io, " "^(indent), "),")
+    return nothing
+end
+
+function _print_type_rule(io::IO, r::TypeRulePair)
+    printstyled(io, r.first; color=COLOR_TYPE)
+    print(io, " => ")
+    printstyled(io, r.second; color=COLOR_RULE)
     return nothing
 end

test/references/show/EpsilonAlpha2Beta1.txt

@@ -1,5 +1,5 @@
 Composite(
-  GlobalTypeRule(  # on all layers
+  GlobalTypeRule(     # on all layers
     Conv => AlphaBetaRule{Float32}(2.0f0, 1.0f0),
     Dense => EpsilonRule{Float32}(1.0f-6),
     Union{typeof(Flux.dropout), AlphaDropout, Dropout} => PassRule(),

test/references/show/EpsilonAlpha2Beta1Flat.txt

@@ -1,11 +1,11 @@
 Composite(
-  GlobalTypeRule(  # on all layers
+  GlobalTypeRule(     # on all layers
     Conv => AlphaBetaRule{Float32}(2.0f0, 1.0f0),
     Dense => EpsilonRule{Float32}(1.0f-6),
     Union{typeof(Flux.dropout), AlphaDropout, Dropout} => PassRule(),
     Union{typeof(flatten), typeof(MLUtils.flatten)} => PassRule(),
   ),
-  FirstLayerTypeRule(  # on first layer
+  FirstLayerTypeRule( # on first layer
     Conv => FlatRule(),
     Dense => FlatRule(),
   ),

test/references/show/EpsilonGammaBox.txt

@@ -1,11 +1,11 @@
 Composite(
-  GlobalTypeRule(  # on all layers
+  GlobalTypeRule(     # on all layers
     Conv => GammaRule{Float32}(0.25f0),
     Dense => EpsilonRule{Float32}(1.0f-6),
     Union{typeof(Flux.dropout), AlphaDropout, Dropout} => PassRule(),
     Union{typeof(flatten), typeof(MLUtils.flatten)} => PassRule(),
   ),
-  FirstLayerTypeRule(  # on first layer
+  FirstLayerTypeRule( # on first layer
     Conv => ZBoxRule{Float32}(-3.0f0, 3.0f0),
   ),
 )

test/references/show/EpsilonPlus.txt

@@ -1,5 +1,5 @@
 Composite(
-  GlobalTypeRule(  # on all layers
+  GlobalTypeRule(     # on all layers
     Conv => ZPlusRule(),
     Dense => EpsilonRule{Float32}(1.0f-6),
     Union{typeof(Flux.dropout), AlphaDropout, Dropout} => PassRule(),

test/references/show/EpsilonPlusFlat.txt

@@ -1,11 +1,11 @@
 Composite(
-  GlobalTypeRule(  # on all layers
+  GlobalTypeRule(     # on all layers
     Conv => ZPlusRule(),
     Dense => EpsilonRule{Float32}(1.0f-6),
     Union{typeof(Flux.dropout), AlphaDropout, Dropout} => PassRule(),
     Union{typeof(flatten), typeof(MLUtils.flatten)} => PassRule(),
   ),
-  FirstLayerTypeRule(  # on first layer
+  FirstLayerTypeRule( # on first layer
     Conv => FlatRule(),
     Dense => FlatRule(),
   ),

test/references/show/composite1.txt

@@ -1,18 +1,18 @@
 Composite(
   GlobalRule: all layers => ZeroRule(),
   GlobalRule: all layers => PassRule(),
-  GlobalTypeRule(  # on all layers
+  GlobalTypeRule(     # on all layers
     Conv => AlphaBetaRule{Float32}(2.0f0, 1.0f0),
     Dense => EpsilonRule{Float32}(1.0f-6),
     Union{GlobalMaxPool, GlobalMeanPool, AdaptiveMaxPool, AdaptiveMeanPool, MaxPool, MeanPool} => EpsilonRule{Float32}(1.0f-6),
   ),
-  FirstNTypeRule(  # on layers 1:7
+  FirstNTypeRule(     # on layers 1:7
     Conv => FlatRule(),
   ),
-  LastNTypeRule(  # on last 3 layers
+  LastNTypeRule(      # on last 3 layers
     Dense => EpsilonRule{Float32}(1.0f-7),
   ),
-  RangeTypeRule(  # on layers 4:10
+  RangeTypeRule(      # on layers 4:10
     Union{GlobalMaxPool, GlobalMeanPool, AdaptiveMaxPool, AdaptiveMeanPool, MaxPool, MeanPool} => EpsilonRule{Float32}(1.0f-5),
   ),
   LayerRule: layer 9 => AlphaBetaRule{Float32}(1.0f0, 0.0f0),

test/references/show/composite2.txt

@@ -3,7 +3,7 @@ Composite(
     Dense => EpsilonRule{Float32}(2.0f-5),
     Conv => EpsilonRule{Float32}(0.0002f0),
   ),
-  FirstLayerTypeRule(  # on first layer
+  FirstLayerTypeRule( # on first layer
     Dense => AlphaBetaRule{Float32}(1.0f0, 0.0f0),
     Conv => AlphaBetaRule{Float32}(2.0f0, 1.0f0),
   ),