diff --git a/src/outputsize.jl b/src/outputsize.jl index 76c58237be..9fd9545b5f 100644 --- a/src/outputsize.jl +++ b/src/outputsize.jl @@ -248,13 +248,13 @@ function _makelazy(ex::Expr) n == 0 && return ex n == 1 && error("@autosize doesn't expect an underscore here: $ex") n == 2 && return :($LazyLayer($(string(ex)), $(_makefun(ex)), nothing)) - n > 2 && return Expr(ex.head, ex.args[1], map(_makelazy, ex.args[2:end])...) + n > 2 && return Expr(ex.head, map(_makelazy, ex.args)...) end _makelazy(x) = x function _underscoredepth(ex::Expr) # Meta.isexpr(ex, :tuple) && :_ in ex.args && return 10 - ex.head in (:call, :kw, :(->), :block) || return 0 + ex.head in (:call, :kw, :(->), :block, :parameters) || return 0 ex.args[1] === :(=>) && ex.args[2] === :_ && return 1 m = maximum(_underscoredepth, ex.args) m == 0 ? 0 : m+1 @@ -279,6 +279,7 @@ is needed to make `@autosize (2,3,4) Dense(_ => 5)` return """ autosizefor(::Type, x::AbstractArray) = size(x, max(1, ndims(x)-1)) autosizefor(::Type{<:Dense}, x::AbstractArray) = size(x, 1) +autosizefor(::Type{<:Embedding}, x::AbstractArray) = size(x, 1) autosizefor(::Type{<:LayerNorm}, x::AbstractArray) = size(x, 1) _replaceunderscore(e, s) = e === :_ ? s : e diff --git a/test/outputsize.jl b/test/outputsize.jl index eec6880dc2..0e5b807a60 100644 --- a/test/outputsize.jl +++ b/test/outputsize.jl @@ -174,12 +174,20 @@ end m = @autosize (2, 3, 4, 5) Dense(_ => 10) # goes by first dim, not 2nd-last @test randn(2, 3, 4, 5) |> m |> size == (10, 3, 4, 5) - + + @test_broken begin # outputsize fails on Embedding + m = @autosize (2, 3, 4, 5) Embedding(_ => 10) # goes by first dim, not 2nd-last + @test randn(2, 3, 4, 5) |> m |> size == (10, 3, 4, 5) + end + m = @autosize (9,) Dense(_ => div(_,2)) @test randn(9) |> m |> size == (4,) m = @autosize (3,) Chain(one = Dense(_ => 4), two = softmax) # needs kw @test randn(3) |> m |> size == (4,) + + m = @autosize (3,) Chain(; one = Dense(_ => 4), two = softmax) # needs parameters + @test randn(3) |> m |> size == (4,) m = @autosize (3, 45) Maxout(() -> Dense(_ => 6, tanh), 2) # needs ->, block @test randn(3, 45) |> m |> size == (6, 45) @@ -222,6 +230,10 @@ end Dense(_ => 10), ) @test randn(Float32, img..., 1, 32) |> m |> size == (10, 32) + + # https://github.com/FluxML/Flux.jl/issues/2086 + m = @autosize (3, 1) Chain(; c = Dense(_ => 2, sigmoid), b = BatchNorm(_, affine=false)) + @test randn(Float32, 3, 32) |> m |> size == (2, 32) end @testset "LazyLayer" begin