Skip to content

Conversation

mcabbott
Copy link
Member

@mcabbott mcabbott commented Aug 29, 2023

This simplifies @compact so that it no longer stores a string for each keyword argument. Like real keywords, they run once, and the resulting NamedTuple is all that remains.

Instead, it proposes to print something like summary(w) for any array. Maybe this needs to be trimmed a bit, as for CuArray etc it can get quite long. Could trim it all the way to 32×26 AbstractMatrix{Float32}. This change should probably be made in Flux anyway (along with other fixes to how show handles arrays) so is just prototyped piratically here for now. [Edit: now https://github.com/FluxML/Flux.jl/pull/2344]

julia> using Flux; using Fluxperimental: @compact

julia> let n = 26, d = 32  # previously the constructor string was printed, although it cannot run again
        Chain(
          @compact(w=randn32(d, n)./=sqrt(n), pow=1+1) do (x, y)
            @views (w[:,x] .+ w[:,y]).^pow
          end,
          Dense(d => 3),
         )
       end
Chain(
  @compact(
    w = 32×26 Matrix{Float32},          # 832 parameters
    pow = 2,
  ) do (x, y) 
      #= REPL[54]:4 =# @views (w[:, x] .+ w[:, y]) .^ pow
  end,
  Dense(32 => 3),                       # 99 parameters
)                   # Total: 3 arrays, 931 parameters, 4.011 KiB.

julia> ans((1,2))
3-element Vector{Float32}:
  0.30514276
 -0.24690717
 -0.26537845

julia> let n = 3, m = 4  # Scale previously triggered the "non-layer" printing path. Note captured m still printed.
         @compact(s = Flux.Scale(max(2n, m))) do x
           s(x ./ m)
         end
       end
@compact(
  s = Scale(6),                         # 12 parameters
) do x 
    s(x ./ m)
end

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant