Add a macro to opt-in to fancy printing, and to everything else #1932

mcabbott · 2022-04-07T00:08:08Z

Originally, `@big_show`

This adds a macro which lets you tell Flux to treat some container type much the same way as it treats Chain, starting the recursive printing walk over all children.

Prompted by #1929. Closes #2044

Metalhead also has show methods for top-level layers which enable the recursive printing, which could be replaced by this macro:

https://github.com/FluxML/Metalhead.jl/blob/aba6fb832093d88dc2d2b4d5b1d2d63a0f21eb9c/src/Metalhead.jl#L53-L57

https://github.com/FluxML/Metalhead.jl/blob/c8f0a88e4d24274c53b62c2b740822aa6a781709/src/utilities.jl#L49-L52

Searching for other uses on github, I see only two:

https://github.com/avik-pal/ExplicitFluxLayers.jl/blob/8e1cff447afda225bc12144ff27ae1370e8fc3da/src/show_layers.jl

https://github.com/maxfreu/SegmentationModels.jl/blob/7bfdbaa438910baf49543f03f2931de765dbd761/src/unet.jl#L99-L112

Later, `@layer`

Now aims to also replace @functor, in addition to handling show.

Discussed at #2028 -- a layer supertype is another way to make opting into pretty printing easier. (We wouldn't do both, so this closes #2028 .)

PR Checklist

Tests are added
Entry in NEWS.md
Documentation, if applicable

src/layers/show.jl

ToucheSir · 2022-04-07T15:42:29Z

Looks reasonable, so if you'd let me indulge in some light bikeshedding how about @show_expanded over @big_show?

mcabbott · 2022-04-07T16:45:06Z

Not at all set on the name.

Another question is whether there should be a matching way to disable recursion. Overloading _show_leaflike(::Diagonal) to make LayerNorm not expand is slightly weird. It's a hack to re-use the fact that layers all of whose children area leaflike aren't expanded, and has side-effects like this:

julia> Chain(Flux.Scale(2), Flux.Scale(2))
Chain(Scale(2), Scale(2))  # 8 parameters

julia> Chain(Flux.Scale(2), Flux.Scale(2), Dense(2=>2))
Chain(
  Scale(2),                             # 4 parameters
  Scale(2),                             # 4 parameters
  Dense(2 => 2),                        # 6 parameters
)                   # Total: 6 arrays, 14 parameters, 440 bytes.

If this needs a public face, then perhaps it should be more like @show_noexpand LayerNorm i.e. a property of the container, not the contents.

CarloLucibello · 2022-06-02T07:28:54Z

show_leaflike(x::MyType) = true allowing to customize on the specific object seems more flexible and it feels ok to me as an interface, I would prefer that over a macro.

We could also merge this PR now since it seems ready and expose that interface later

mcabbott · 2022-08-23T20:51:41Z

6370374 upgrades this to a replacement for @functor. The idea is to have just one magic macro do several things for us, maybe that's less confusing. It would also allow any back-end changes later, once Flux owns the macro. Discussed at #2028 (as an alternative to the supertype idea there).

Still WIP, but roughly works, I think.

ToucheSir · 2022-08-24T02:26:19Z

src/layers/macro.jl

+* In fact you can provide an arbitrary keyword with this syntax, and it will
+  overload this function alla `trainable`... that might be a terrible idea.


As a start, we could limit the allowable keyword list.

Yes. Nice to pick a notation that can grow if desired though.

ToucheSir · 2022-08-24T02:32:09Z

src/layers/macro.jl

+    @layer Dense
+    @layer :expand Chain
+    @layer BatchNorm trainable=(β,γ)
+    @layer Struct functor=(α,β) trainable=(β,)


While we're at it, perhaps functor could be aliased to something less obscure? params might be a name conflict but is the most familiar (perhaps parameters? Too cheeky?). buffers is close to PyTorch but still kind of obscure.

Why not leaves or children?

Could do. I guess I wanted to match the function name, and if this one is obscure, that's OK, you should almost never use it. children does match a function name.

children also has the benefit of having the same signature as trainable: () -> Union{Tuple,NamedTuple,Array,...} instead of () -> Tuple{Union{Tuple,NamedTuple,Array,...}, Function} for functor.

Latest version uses children as the keyword.

mcabbott · 2022-08-25T15:44:03Z

src/layers/macro.jl

+function _check_new_macro(x::T) where T
+  Functors.isleaf(x) && return
+  Base.depwarn("This type should probably now use `Flux.@layer` instead of `@functor`: $T", Symbol("@functor"))
+end


I toned this warning down from @warn. I believe this makes the PR non-breaking, could be in a 0.13 release. Possibly 0.14 should always warn?

A related question. Currently many people use Flux.@functor. We could define such a macro as a pass-through to @layer, which will print a depwarn once on the macro call. Should we?

src/layers/macro.jl

mcabbott · 2022-08-25T15:47:16Z

src/layers/macro.jl

+   else
+     remake(nt) = Base.typename(T).wrapper(map(f -> f in which ? getfield(nt, f) : getfield(x, f), fieldnames(T))...)
+     NamedTuple{which}(map(s -> getfield(x, s), which)), remake
+   end


The path when you specify children isn't as fast. Maybe this doesn't matter for now, since no layer within Flux does so.

This may be why Functors.jl uses @eval. A generated block seems fine though.

For sure it can be done with a @generated block. I was just lazy to actually write and check and benchmark it.

Can this be done with @generated? It seems you can't define closures in generated functions, and we need to close over x.

I guess that's true. But perhaps it can instead return a struct like Remake{(:some, :fields)}(x, constructor) which gets some fields from x and some from what it's called on. But I didn't try.

I'm also horrified by the number of ways to customise recursive behaviour we have, and I wonder if this one should just be forbidden.

Clever, let me try that.

Wait isn't returning a Remake struct the same as returning a closure (in terms of performance not whether it can be in a generated function)? I thought the benefit here was that you can manually write out the iteration accessing the fields.

Yes I would guess so. If necc. its call could itself be a generated thing, but quite likely some ntuple / map would just compile down.

Not precisely sure what list of fields this thing ought to be passed, the ones you get or the ones you don't or some mix such that merge sorts it out?

mcabbott · 2022-08-25T15:48:18Z

src/layers/normalise.jl

-@functor BatchNorm
-trainable(bn::BatchNorm) = hasaffine(bn) ? (β = bn.β, γ = bn.γ) : (;)
+@layer BatchNorm trainable=(β,γ)


Normalisation layers all have permanent trainable, but the fields are nothing when not in use, should be safe?

Yes. Really, the affine params should be wrapped up in a Scale sublayer like LayerNorm has, but that's a PR for another day.

src/layers/recurrent.jl

mcabbott · 2022-08-25T15:54:39Z

src/layers/show.jl

@@ -1,22 +1,30 @@
+@nospecialize  # just for this file, for startup time


Sometimes show is visibly slow on first loading. This doesn't really help, a test with @time show(stdout, MIME"text/plain"(), model); goes from 2.28s to 2.15s.

mcabbott · 2022-08-25T15:56:02Z

docs/src/models/basics.md

@@ -216,7 +216,7 @@ m(5) # => 26
 Flux provides a set of helpers for custom layers, which you can enable by calling

 ```julia
-Flux.@functor Affine
+Flux.@layer Affine


"set of helpers for custom layers" is more accurate after than before. Where else would docs need changing?

https://github.com/FluxML/Flux.jl/blob/master/docs/src/models/advanced.md is the big one, off the top of my head.

I had a go at editing that. It's quite a crazy scheme we have going, where trainable, functor, adapt define 3 different recursions. Do we really need to ways to customise functor?

ToucheSir · 2022-08-26T15:11:41Z

Doctest failure looks real, but I'm not sure why it's only showing up now.

mcabbott · 2022-08-26T20:11:21Z

Doc failure seems to be that this doesn't work on 1.6, where the first argument is a vector:

julia> Dense([3.3;;], false)
Dense(1 => 1; bias=false)  # 1 parameters

Downstream tests didn't run, not sure why.

mcabbott · 2022-08-26T23:44:45Z

Failures look unrelated:

GeometricFlux.jl I see only create_bias failures, Make create_bias a public API? #2049
FastAI.jl has failures like "Expression: isempty(stderr_content) [1105] Evaluated: isempty("┌ Warning: ignore(f) is deprecated, use ChainRulesCore.ignore_derivatives(f) instead.\n│ caller ="
AtomicGraphNets.jl has failures like "MethodError: no method matching needs_concrete_A(::Nothing) Closest candidates are: needs_concrete_A(::LinearSolve.AbstractFactorization) at ~/.julia/packages/LinearSolve/AoYJI/src/LinearSolve.jl:34"
OperatorLearning.jl has create_bias only

ToucheSir · 2022-08-27T02:21:38Z

So just the doctest fixup for 1.6?

fixup tidy up, add NEWS review suggestions macro docstring, incl. hcat(3.3)

codecov · 2024-02-29T14:56:51Z

Codecov Report

Attention: Patch coverage is 70.19231% with 31 lines in your changes are missing coverage. Please review.

Project coverage is 73.91%. Comparing base (20d516b) to head (63e10f6).

Files	Patch %	Lines
src/layers/attention.jl	9.52%	19 Missing ⚠️
src/layers/macro.jl	81.66%	11 Missing ⚠️
src/layers/normalise.jl	0.00%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #1932       +/-   ##
===========================================
+ Coverage   42.55%   73.91%   +31.35%     
===========================================
  Files          31       32        +1     
  Lines        1786     1909      +123     
===========================================
+ Hits          760     1411      +651     
+ Misses       1026      498      -528

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

mcabbott · 2024-02-29T15:49:00Z

Rebased in 2024. We should just do this, right?

The biggest questions is: Should this support children at all? Or can we simply forbid that... do we really need 4 ways to customise things, alla #2385? Maybe if you want to hack things with Functors.jl directly you are welcome, but Flux.jl can pretend this doesn't exist.

darsnack · 2024-02-29T18:33:52Z

I agree. Given that you almost never want to set children, having to support and explain children vs trainable is just confusing for users. Anyone wanting to restrict children probably needs to take a minute to understand @functor first.

mcabbott · 2024-03-02T20:10:41Z

OK, children is gone.

This should be ready to merge, I think. As always there may be further cleaning up of the docs possible.

ToucheSir

Let's do this.

…ML#1932)

mcabbott mentioned this pull request Apr 7, 2022

Chain ignores Base.show function of custom layer #1929

Closed

mcabbott commented Apr 7, 2022

View reviewed changes

src/layers/show.jl Outdated Show resolved Hide resolved

CarloLucibello mentioned this pull request Jun 2, 2022

add show methods for WithGraph and GNNChain JuliaGraphs/GraphNeuralNetworks.jl#178

Closed

mcabbott mentioned this pull request Jul 27, 2022

RFC: add a supertype to layers #2028

Closed

mcabbott mentioned this pull request Aug 22, 2022

Macro to display model struct the way Flux does #2044

Closed

mcabbott force-pushed the showmacro branch from 320cad5 to 6370374 Compare August 23, 2022 20:49

ToucheSir reviewed Aug 24, 2022

View reviewed changes

mcabbott force-pushed the showmacro branch 2 times, most recently from 85252bf to 8981283 Compare August 24, 2022 21:21

mcabbott changed the title ~~Add a macro to opt-in to fancy printing~~ Add a macro to opt-in to fancy printing, and to everything else Aug 24, 2022

mcabbott marked this pull request as ready for review August 25, 2022 14:57

mcabbott commented Aug 25, 2022

View reviewed changes

ToucheSir added the run downstream test label Aug 26, 2022

ToucheSir closed this Aug 26, 2022

ToucheSir reopened this Aug 26, 2022

mcabbott added this to the v0.14 milestone Aug 28, 2022

mcabbott force-pushed the showmacro branch 2 times, most recently from d891f6d to 55ebc57 Compare August 29, 2022 16:14

theabhirath mentioned this pull request Sep 7, 2022

Abstract type for the models FluxML/Metalhead.jl#183

Open

mcabbott added the discussion label Sep 16, 2022

mcabbott mentioned this pull request Oct 11, 2022

Separate walks out from fmap and add #39 to fcollect FluxML/Functors.jl#43

Merged

mcabbott force-pushed the showmacro branch from 55ebc57 to f2a2cbb Compare October 11, 2022 15:27

mcabbott mentioned this pull request Nov 19, 2022

Re-organise "built-in layers" section #2112

Merged

mcabbott force-pushed the showmacro branch from f2a2cbb to 445db4b Compare November 20, 2022 17:53

mcabbott mentioned this pull request Nov 25, 2022

functor by default FluxML/Functors.jl#51

Merged

3 tasks

CarloLucibello modified the milestones: v0.14, v0.15 Jul 10, 2023

mcabbott added 2 commits February 29, 2024 09:39

add big_show macro

7ca3449

upgrade to at-layer macro, replaces at-functor

3ab088b

fixup tidy up, add NEWS review suggestions macro docstring, incl. hcat(3.3)

mcabbott force-pushed the showmacro branch from 445db4b to 3ab088b Compare February 29, 2024 14:56

minor fixes after rebase

35ade0a

mcabbott added 4 commits March 1, 2024 20:45

layer MultiHeadAttention, and show methods for this

29e0d68

remove children=(α,β) keyword

2be9099

fixup attention

fe72081

rm note

63e10f6

ToucheSir approved these changes Mar 5, 2024

View reviewed changes

mcabbott merged commit 5e80211 into FluxML:master Mar 5, 2024
18 of 20 checks passed

mcabbott deleted the showmacro branch March 5, 2024 20:38

This was referenced Mar 6, 2024

doc changes re at-functor and at-layer #2390

Merged

docs: improve freezing docs #2385

Closed

isentropic pushed a commit to isentropic/Flux.jl that referenced this pull request Mar 13, 2024

Add a macro to opt-in to fancy printing, and to everything else (Flux…

97fdcd1

…ML#1932)

		* In fact you can provide an arbitrary keyword with this syntax, and it will
		overload this function alla `trainable`... that might be a terrible idea.

		@@ -1,22 +1,30 @@
		@nospecialize # just for this file, for startup time

Add a macro to opt-in to fancy printing, and to everything else #1932

Add a macro to opt-in to fancy printing, and to everything else #1932

Conversation

mcabbott commented Apr 7, 2022 • edited Loading

Originally, @big_show

Later, @layer

PR Checklist

ToucheSir commented Apr 7, 2022

mcabbott commented Apr 7, 2022

CarloLucibello commented Jun 2, 2022

mcabbott commented Aug 23, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ToucheSir Aug 24, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ToucheSir Aug 26, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ToucheSir commented Aug 26, 2022

mcabbott commented Aug 26, 2022 • edited Loading

mcabbott commented Aug 26, 2022

ToucheSir commented Aug 27, 2022

codecov bot commented Feb 29, 2024 • edited Loading

Codecov Report

mcabbott commented Feb 29, 2024

darsnack commented Feb 29, 2024

mcabbott commented Mar 2, 2024

ToucheSir left a comment

Choose a reason for hiding this comment

mcabbott commented Apr 7, 2022 •

edited

Loading

Originally, `@big_show`

Later, `@layer`

mcabbott commented Aug 23, 2022 •

edited

Loading

ToucheSir Aug 24, 2022 •

edited

Loading

ToucheSir Aug 26, 2022 •

edited

Loading

mcabbott commented Aug 26, 2022 •

edited

Loading

codecov bot commented Feb 29, 2024 •

edited

Loading