From 43fee12c369d752eba9928a545e63abb454bf956 Mon Sep 17 00:00:00 2001
From: Michael Abbott <32575566+mcabbott@users.noreply.github.com>
Date: Sun, 16 Oct 2022 09:30:49 -0400
Subject: [PATCH 1/2] fix 2086

---
 src/outputsize.jl  | 4 ++--
 test/outputsize.jl | 7 +++++++
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/outputsize.jl b/src/outputsize.jl
index 76c58237be..952de5487e 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
diff --git a/test/outputsize.jl b/test/outputsize.jl
index eec6880dc2..af88d8b038 100644
--- a/test/outputsize.jl
+++ b/test/outputsize.jl
@@ -180,6 +180,9 @@ end
 
   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 +225,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

From 9c47f72e0da0aa0468117782d451659ed8c8fb3d Mon Sep 17 00:00:00 2001
From: Michael Abbott <32575566+mcabbott@users.noreply.github.com>
Date: Sun, 16 Oct 2022 09:31:11 -0400
Subject: [PATCH 2/2] Embedding, but not yet

---
 src/outputsize.jl  | 1 +
 test/outputsize.jl | 7 ++++++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/outputsize.jl b/src/outputsize.jl
index 952de5487e..9fd9545b5f 100644
--- a/src/outputsize.jl
+++ b/src/outputsize.jl
@@ -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 af88d8b038..0e5b807a60 100644
--- a/test/outputsize.jl
+++ b/test/outputsize.jl
@@ -174,7 +174,12 @@ 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,)