From 662ef5f8e98e4b73b3e7137c4cc1d8e8856fe35f Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Sat, 16 Jan 2021 21:45:13 +0100
Subject: [PATCH 01/19] add Upsample and PixelShuffle layers

---
 docs/src/models/layers.md |  7 +++++++
 docs/src/models/nnlib.md  |  7 +++++++
 src/Flux.jl               |  2 ++
 src/layers/upsample.jl    | 34 ++++++++++++++++++++++++++++++++++
 test/layers/upsample.jl   | 13 +++++++++++++
 test/runtests.jl          |  1 +
 6 files changed, 64 insertions(+)
 create mode 100644 src/layers/upsample.jl
 create mode 100644 test/layers/upsample.jl

diff --git a/docs/src/models/layers.md b/docs/src/models/layers.md
index 34dc1b57f5..ed6027f5ca 100644
--- a/docs/src/models/layers.md
+++ b/docs/src/models/layers.md
@@ -29,6 +29,13 @@ Flux.convfilter
 Flux.depthwiseconvfilter
 ```
 
+## Upsampling Layers
+
+```@docs
+Upsample
+PixelShuffle
+```
+
 ## Recurrent Layers
 
 Much like the core layers above, but can be used to process sequence data (as well as other kinds of structured data).
diff --git a/docs/src/models/nnlib.md b/docs/src/models/nnlib.md
index 7ede2682ac..8216dcfd85 100644
--- a/docs/src/models/nnlib.md
+++ b/docs/src/models/nnlib.md
@@ -51,6 +51,13 @@ NNlib.conv
 NNlib.depthwiseconv
 ```
 
+## Upsampling
+
+```@docs
+NNlib.upsample_bilinear
+NNlib.pixel_shuffle
+```
+
 ## Batched Operations
 
 ```@docs
diff --git a/src/Flux.jl b/src/Flux.jl
index ff067803be..5e6776d601 100644
--- a/src/Flux.jl
+++ b/src/Flux.jl
@@ -16,6 +16,7 @@ export Chain, Dense, Maxout, SkipConnection, Parallel, flatten,
        SamePad, Conv, CrossCor, ConvTranspose, DepthwiseConv,
        AdaptiveMaxPool, AdaptiveMeanPool, GlobalMaxPool, GlobalMeanPool, MaxPool, MeanPool,
        Dropout, AlphaDropout, LayerNorm, BatchNorm, InstanceNorm, GroupNorm,
+       Upsample, PixelShuffle,
        params, fmap, cpu, gpu, f32, f64,
        testmode!, trainmode!
 
@@ -42,6 +43,7 @@ include("layers/basic.jl")
 include("layers/conv.jl")
 include("layers/recurrent.jl")
 include("layers/normalise.jl")
+include("layers/upsample.jl")
 
 include("outputsize.jl")
 
diff --git a/src/layers/upsample.jl b/src/layers/upsample.jl
new file mode 100644
index 0000000000..c9bcac6397
--- /dev/null
+++ b/src/layers/upsample.jl
@@ -0,0 +1,34 @@
+"""
+    Upsample(; scale_factor, mode)
+
+An upsampling layer.
+`scale_factor`: a tuple of ints. 
+`mode`: only `:bilinear` currently supported. 
+
+See [`NNlib.upsample_bilinear`](@ref).
+"""
+struct Upsample{Mode, N, T} 
+  scale_factor::NTuple{N, T}
+end
+
+function Upsample(; scale_factor::NTuple{N, T}, mode::Symbol) where {N,T}
+  @assert mode in [:bilinear] "Not supported mode '$mode'"
+  @assert N == 2  "Only 2d upsampling currently supported"
+  Upsample{Val{mode}, N, T}(scale_factor)
+end
+
+(m::Upsample{Val{:bilinear}})(x) = NNlib.upsample_bilinear(x, m.scale_factor)
+
+
+"""
+    PixelShuffle(r::Int)
+
+Pixel shuffling layer with upscale factor `r`.
+ 
+See [`NNlib.pixel_shuffle`](@ref).
+"""
+struct PixelShuffle 
+  r::Int
+end
+
+(m::PixelShuffle)(x) = NNlib.pixel_shuffle(x, m.r)
diff --git a/test/layers/upsample.jl b/test/layers/upsample.jl
new file mode 100644
index 0000000000..df4d51e232
--- /dev/null
+++ b/test/layers/upsample.jl
@@ -0,0 +1,13 @@
+@testset "bilinear upsample" begin
+  m = Upsample(scale_factor=(2, 3), mode=:bilinear)
+  y = m((rand(Float32, 3, 4, 2, 3)))
+  @test y isa Array{Float32, 4} 
+  @test size(y) == (6, 12, 2, 3)
+end
+
+@testset "PixelShuffle" begin
+  m = PixelShuffle(3)
+  y = m((rand(Float32, 3, 4, 18, 3)))
+  @test y isa Array{Float32, 4} 
+  @test size(y) == (9, 12, 2, 3)
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 0633d14e7f..a40433d0f1 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -34,6 +34,7 @@ end
   include("layers/stateless.jl")
   include("layers/recurrent.jl")
   include("layers/conv.jl")
+  include("layers/upsample.jl")
 end
 
 @testset "outputsize" begin

From 3db3dad6342ccd632c1a859c169f615549562d52 Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Sun, 17 Jan 2021 00:18:19 +0100
Subject: [PATCH 02/19] scale; add news; docs

---
 NEWS.md                 |  2 ++
 src/layers/upsample.jl  | 32 +++++++++++++++++++++++---------
 test/layers/upsample.jl |  2 +-
 3 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/NEWS.md b/NEWS.md
index db7853995e..6f77ede84f 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -12,6 +12,8 @@
 * Moved GPU CI to use buildkite instead of GitLab
 * New [`Parallel` layer](https://github.com/FluxML/Flux.jl/pull/1462) adds inception module-like building blocks.
 * Feature additions and bug fixes for BatchNorm, LayerNorm, InstanceNorm, and GroupNorm [normalization layers](https://github.com/FluxML/Flux.jl/pull/1397)
+* Added [Upsample and PixelShuffle layers](https://github.com/FluxML/Flux.jl/pull/1468)
+* Other new features and bug fixes (see GitHub releases page)
 
 ## v0.11.2
 
diff --git a/src/layers/upsample.jl b/src/layers/upsample.jl
index c9bcac6397..beda3cf368 100644
--- a/src/layers/upsample.jl
+++ b/src/layers/upsample.jl
@@ -1,23 +1,37 @@
 """
-    Upsample(; scale_factor, mode)
+    Upsample(; scale, mode)
 
-An upsampling layer.
-`scale_factor`: a tuple of ints. 
-`mode`: only `:bilinear` currently supported. 
+An upsampling layer. `scale` is a tuple determining 
+the output rescaling factor along each spatial dimension.
+The only currently supported upsampling `mode` is 
+`:bilinear`.
 
 See [`NNlib.upsample_bilinear`](@ref).
+
+# Examples
+
+```juliarepl
+julia> m = Upsample(scale=(2,3), mode=:bilinear)
+Upsample{Val{:bilinear},2,Int64}((2, 3))
+
+julia> m(ones(1,1,1,1))
+2×3×1×1 Array{Float64,4}:
+[:, :, 1, 1] =
+ 1.0  1.0  1.0
+ 1.0  1.0  1.0
+```
 """
 struct Upsample{Mode, N, T} 
-  scale_factor::NTuple{N, T}
+  scale::NTuple{N, T}
 end
 
-function Upsample(; scale_factor::NTuple{N, T}, mode::Symbol) where {N,T}
-  @assert mode in [:bilinear] "Not supported mode '$mode'"
+function Upsample(; scale::NTuple{N, T}, mode::Symbol) where {N,T}
+  @assert mode in [:bilinear] "Unsupported mode '$mode'"
   @assert N == 2  "Only 2d upsampling currently supported"
-  Upsample{Val{mode}, N, T}(scale_factor)
+  Upsample{Val{mode}, N, T}(scale)
 end
 
-(m::Upsample{Val{:bilinear}})(x) = NNlib.upsample_bilinear(x, m.scale_factor)
+(m::Upsample{Val{:bilinear}})(x) = NNlib.upsample_bilinear(x, m.scale)
 
 
 """
diff --git a/test/layers/upsample.jl b/test/layers/upsample.jl
index df4d51e232..9f6349358e 100644
--- a/test/layers/upsample.jl
+++ b/test/layers/upsample.jl
@@ -1,5 +1,5 @@
 @testset "bilinear upsample" begin
-  m = Upsample(scale_factor=(2, 3), mode=:bilinear)
+  m = Upsample(scale=(2, 3), mode=:bilinear)
   y = m((rand(Float32, 3, 4, 2, 3)))
   @test y isa Array{Float32, 4} 
   @test size(y) == (6, 12, 2, 3)

From 39b892de81484a24276140ab5cfe021cb700019e Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@unibocconi.it>
Date: Sun, 17 Jan 2021 07:54:43 +0100
Subject: [PATCH 03/19] Update src/layers/upsample.jl

Co-authored-by: Michael Abbott <32575566+mcabbott@users.noreply.github.com>
---
 src/layers/upsample.jl | 21 ++++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/src/layers/upsample.jl b/src/layers/upsample.jl
index beda3cf368..684d8b5bc7 100644
--- a/src/layers/upsample.jl
+++ b/src/layers/upsample.jl
@@ -21,17 +21,24 @@ julia> m(ones(1,1,1,1))
  1.0  1.0  1.0
 ```
 """
-struct Upsample{Mode, N, T} 
-  scale::NTuple{N, T}
+struct Upsample{Mode,T}
+  scale::T
 end
 
-function Upsample(; scale::NTuple{N, T}, mode::Symbol) where {N,T}
-  @assert mode in [:bilinear] "Unsupported mode '$mode'"
-  @assert N == 2  "Only 2d upsampling currently supported"
-  Upsample{Val{mode}, N, T}(scale)
+function Upsample(; scale, mode::Symbol=:bilinear)
+  mode in [:bilinear] || throw(ArgumentError("only bilinear sampling is supported"))
+  scale isa Integer || scale isa NTuple{2,<:Integer} || 
+    throw(ArgumentError("only two-dimensional scaling by an integer factor is supported"))
+  Upsample{mode, typeof(scale)}(scale)
 end
 
-(m::Upsample{Val{:bilinear}})(x) = NNlib.upsample_bilinear(x, m.scale)
+(m::Upsample{:bilinear})(x) = NNlib.upsample_bilinear(x, m.scale isa Tuple ? m.scale : (m.scale, m.scale))
+
+function Base.show(io::IO, u::Upsample2{mode}) where {mode}
+  print(io, "Upsample(scale=", u.scale)
+  mode == :bilinear || print(io, ", mode=", mode)
+  print(io, ")")
+end
 
 
 """

From 96ab813158baf3656c24c24b1a12a9e482a3eda1 Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Sun, 17 Jan 2021 08:15:32 +0100
Subject: [PATCH 04/19] small fixes

---
 src/layers/upsample.jl | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/src/layers/upsample.jl b/src/layers/upsample.jl
index 684d8b5bc7..de9745ee03 100644
--- a/src/layers/upsample.jl
+++ b/src/layers/upsample.jl
@@ -1,8 +1,8 @@
 """
     Upsample(; scale, mode)
 
-An upsampling layer. `scale` is a tuple determining 
-the output rescaling factor along each spatial dimension.
+An upsampling layer. `scale` is an integer or a tuple of integerers 
+representing  the output rescaling factor along each spatial dimension.
 The only currently supported upsampling `mode` is 
 `:bilinear`.
 
@@ -12,7 +12,7 @@ See [`NNlib.upsample_bilinear`](@ref).
 
 ```juliarepl
 julia> m = Upsample(scale=(2,3), mode=:bilinear)
-Upsample{Val{:bilinear},2,Int64}((2, 3))
+Upsample(scale=(2, 3), mode=bilinear)
 
 julia> m(ones(1,1,1,1))
 2×3×1×1 Array{Float64,4}:
@@ -25,18 +25,19 @@ struct Upsample{Mode,T}
   scale::T
 end
 
-function Upsample(; scale, mode::Symbol=:bilinear)
-  mode in [:bilinear] || throw(ArgumentError("only bilinear sampling is supported"))
-  scale isa Integer || scale isa NTuple{2,<:Integer} || 
-    throw(ArgumentError("only two-dimensional scaling by an integer factor is supported"))
-  Upsample{mode, typeof(scale)}(scale)
+function Upsample(; scale::Union{Int, NTuple{N,Int}}, mode::Symbol) where N
+  mode in [:linear, :bilinear] || 
+    throw(ArgumentError("only sampling is currently supported"))
+  scale isa Int  || N == 2 ||
+    throw(ArgumentError("only two-dimensional scaling  is supported"))
+  return Upsample{mode, typeof(scale)}(scale)
 end
 
 (m::Upsample{:bilinear})(x) = NNlib.upsample_bilinear(x, m.scale isa Tuple ? m.scale : (m.scale, m.scale))
 
-function Base.show(io::IO, u::Upsample2{mode}) where {mode}
+function Base.show(io::IO, u::Upsample{mode}) where {mode}
   print(io, "Upsample(scale=", u.scale)
-  mode == :bilinear || print(io, ", mode=", mode)
+  print(io, ", mode=:", mode)
   print(io, ")")
 end
 

From d3c9cde5848dccf974a69d0375163b34a621bb53 Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Sun, 17 Jan 2021 08:17:14 +0100
Subject: [PATCH 05/19] fix

---
 src/layers/upsample.jl | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/layers/upsample.jl b/src/layers/upsample.jl
index de9745ee03..ce29e7a8e4 100644
--- a/src/layers/upsample.jl
+++ b/src/layers/upsample.jl
@@ -26,10 +26,10 @@ struct Upsample{Mode,T}
 end
 
 function Upsample(; scale::Union{Int, NTuple{N,Int}}, mode::Symbol) where N
-  mode in [:linear, :bilinear] || 
-    throw(ArgumentError("only sampling is currently supported"))
+  mode in [:bilinear] || 
+    throw(ArgumentError("only `:bilinear` upsampling is currently supported"))
   scale isa Int  || N == 2 ||
-    throw(ArgumentError("only two-dimensional scaling  is supported"))
+    throw(ArgumentError("only two-dimensional upsampling  is supported"))
   return Upsample{mode, typeof(scale)}(scale)
 end
 

From d50df60b66f84d986ee0acee5e6b05e346be6b76 Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Wed, 20 Jan 2021 14:40:51 +0100
Subject: [PATCH 06/19] add neirest

---
 Manifest.toml            |  8 ++++----
 docs/src/models/nnlib.md |  1 +
 src/layers/upsample.jl   | 32 +++++++++++++++++++++-----------
 test/layers/upsample.jl  | 40 +++++++++++++++++++++++++++++++++++++---
 4 files changed, 63 insertions(+), 18 deletions(-)

diff --git a/Manifest.toml b/Manifest.toml
index a22b079f19..e3b463a9a2 100644
--- a/Manifest.toml
+++ b/Manifest.toml
@@ -93,9 +93,9 @@ uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
 version = "0.3.4+0"
 
 [[DataAPI]]
-git-tree-sha1 = "6d64b28d291cb94a0d84e6e41081fb081e7f717f"
+git-tree-sha1 = "8ab70b4de35bb3b8cc19654f6b893cf5164f8ee8"
 uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
-version = "1.5.0"
+version = "1.5.1"
 
 [[DataStructures]]
 deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
@@ -236,9 +236,9 @@ uuid = "a63ad114-7e13-5084-954f-fe012c677804"
 
 [[NNlib]]
 deps = ["ChainRulesCore", "Compat", "LinearAlgebra", "Pkg", "Requires", "Statistics"]
-git-tree-sha1 = "573cc0d31f9697b9d2b060130a7a3c05a4f36b78"
+git-tree-sha1 = "e031936930ecb559a92992706f8acb02dc81fed2"
 uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
-version = "0.7.12"
+version = "0.7.13"
 
 [[NaNMath]]
 git-tree-sha1 = "bfe47e760d60b82b66b61d2d44128b62e3a369fb"
diff --git a/docs/src/models/nnlib.md b/docs/src/models/nnlib.md
index 8216dcfd85..d60cc9ea52 100644
--- a/docs/src/models/nnlib.md
+++ b/docs/src/models/nnlib.md
@@ -54,6 +54,7 @@ NNlib.depthwiseconv
 ## Upsampling
 
 ```@docs
+NNlib.upsample_nearest
 NNlib.upsample_bilinear
 NNlib.pixel_shuffle
 ```
diff --git a/src/layers/upsample.jl b/src/layers/upsample.jl
index ce29e7a8e4..9b9040243f 100644
--- a/src/layers/upsample.jl
+++ b/src/layers/upsample.jl
@@ -1,12 +1,14 @@
 """
-    Upsample(; scale, mode)
+  Upsample(scale; mode=:nearest)  
+  Upsample(; scale, mode=:nearest)
 
-An upsampling layer. `scale` is an integer or a tuple of integerers 
+An upsampling layer. `scale` is a number or a tuple of numbers 
 representing  the output rescaling factor along each spatial dimension.
-The only currently supported upsampling `mode` is 
-`:bilinear`.
+Currently supported upsampling `mode`s are:
+  - `:nearest`
+  - `:bilinear`
 
-See [`NNlib.upsample_bilinear`](@ref).
+See [`NNlib.upsample_nearest`](@ref), [`NNlib.upsample_bilinear`](@ref).
 
 # Examples
 
@@ -25,15 +27,23 @@ struct Upsample{Mode,T}
   scale::T
 end
 
-function Upsample(; scale::Union{Int, NTuple{N,Int}}, mode::Symbol) where N
-  mode in [:bilinear] || 
-    throw(ArgumentError("only `:bilinear` upsampling is currently supported"))
-  scale isa Int  || N == 2 ||
-    throw(ArgumentError("only two-dimensional upsampling  is supported"))
+Upsample(; scale, mode::Symbol=:nearest) = Upsample(scale; mode)
+
+function Upsample(scale; mode::Symbol=:nearest)
+  mode in [:nearest, :bilinear] || 
+    throw(ArgumentError("`:$mode` mode is not supported."))
   return Upsample{mode, typeof(scale)}(scale)
 end
 
-(m::Upsample{:bilinear})(x) = NNlib.upsample_bilinear(x, m.scale isa Tuple ? m.scale : (m.scale, m.scale))
+(m::Upsample{:nearest,<:Number})(x::AbstractArray{T,N}) where {T,N} =
+  NNlib.upsample_nearest(x, ntuple(_ -> m.scale, N-2))
+(m::Upsample{:nearest,<:Tuple})(x::AbstractArray) =
+  NNlib.upsample_nearest(x, m.scale)
+
+(m::Upsample{:bilinear,<:Number})(x::AbstractArray{T,N}) where {T,N} =
+  NNlib.upsample_bilinear(x, ntuple(_ -> m.scale, N-2))
+(m::Upsample{:bilinear,<:Tuple})(x::AbstractArray) =
+  NNlib.upsample_bilinear(x, m.scale)
 
 function Base.show(io::IO, u::Upsample{mode}) where {mode}
   print(io, "Upsample(scale=", u.scale)
diff --git a/test/layers/upsample.jl b/test/layers/upsample.jl
index 9f6349358e..936767dcc4 100644
--- a/test/layers/upsample.jl
+++ b/test/layers/upsample.jl
@@ -1,13 +1,47 @@
-@testset "bilinear upsample" begin
+@testset "upsample bilinear" begin
   m = Upsample(scale=(2, 3), mode=:bilinear)
-  y = m((rand(Float32, 3, 4, 2, 3)))
+  x = rand(Float32, 3, 4, 2, 3)
+  y = m(x)
   @test y isa Array{Float32, 4} 
   @test size(y) == (6, 12, 2, 3)
+
+  m = Upsample(3, mode=:bilinear)
+  x = rand(Float32, 3, 4, 2, 3)
+  y = m(x)
+  @test y isa Array{Float32, 4} 
+  @test size(y) == (9, 12, 2, 3)
+end
+
+@testset "upsample nearest" begin
+  x = rand(Float32, 3, 4, 2, 3)
+  m = Upsample(scale=(2, 3), mode=:nearest)
+  y = m(x)
+  @test y isa Array{Float32, 4} 
+  @test size(y) == (6, 12, 2, 3)
+  
+  x = rand(Float32, 3, 4, 2, 3)
+  m = Upsample(2, mode=:nearest)
+  y = m(x)
+  @test y isa Array{Float32, 4} 
+  @test size(y) == (6, 8, 2, 3)
+
+  x = rand(Float32, 3, 2, 3)
+  m = Upsample(2, mode=:nearest)
+  y = m(x)
+  @test y isa Array{Float32, 3} 
+  @test size(y) == (6, 2, 3)
 end
 
 @testset "PixelShuffle" begin
+  m = PixelShuffle(2)
+  x = rand(Float32, 3, 18, 3)
+  y = m(x)
+  @test y isa Array{Float32, 3} 
+  @test size(y) == (6, 9, 3)
+
   m = PixelShuffle(3)
-  y = m((rand(Float32, 3, 4, 18, 3)))
+  x = rand(Float32, 3, 4, 18, 3)
+  y = m(x)
   @test y isa Array{Float32, 4} 
   @test size(y) == (9, 12, 2, 3)
 end

From 320d9a8c13eca1ddf1099174ca829547a6f71d0e Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Thu, 21 Jan 2021 02:43:52 +0100
Subject: [PATCH 07/19] add gpu tests

---
 src/layers/upsample.jl  |  2 +-
 test/cuda/layers.jl     | 25 ++++++++++++++++++++++++-
 test/layers/upsample.jl | 12 ++++++------
 3 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/src/layers/upsample.jl b/src/layers/upsample.jl
index 9b9040243f..52b9b660f1 100644
--- a/src/layers/upsample.jl
+++ b/src/layers/upsample.jl
@@ -31,7 +31,7 @@ Upsample(; scale, mode::Symbol=:nearest) = Upsample(scale; mode)
 
 function Upsample(scale; mode::Symbol=:nearest)
   mode in [:nearest, :bilinear] || 
-    throw(ArgumentError("`:$mode` mode is not supported."))
+    throw(ArgumentError("mode=:$mode is not supported, allowed values are :nearest, :bilinear"))
   return Upsample{mode, typeof(scale)}(scale)
 end
 
diff --git a/test/cuda/layers.jl b/test/cuda/layers.jl
index 8057aa384d..3c44f5cab3 100644
--- a/test/cuda/layers.jl
+++ b/test/cuda/layers.jl
@@ -168,4 +168,27 @@ end
   @test sum(l(ip)) ≈ 0.f0  
   gs = gradient(() -> sum(l(ip)), Flux.params(l))
   @test l.b ∉ gs.params 
-end
\ No newline at end of file
+end
+
+@testset "upsample bilinear" begin
+  m = Upsample(3, mode=:bilinear)
+  x = rand(Float32, 3, 4, 2, 3)
+  gy = m(x |> gpu)
+  @test gy isa CuArray{Float32, 4}
+  @test collect(gy) ≈ m(x)
+end
+
+@testset "upsample nearest" begin
+  x = rand(Float32, 3, 4, 2, 3)
+  gy = m(x |> gpu)
+  @test gy isa CuArray{Float32, 4}
+  @test collect(gy) ≈ m(x)
+end
+
+@testset "PixelShuffle" begin
+  m = PixelShuffle(3)
+  x = rand(Float32, 3, 4, 18, 3)
+  gy = m(x |> gpu)
+  @test gy isa CuArray{Float32, 4}
+  @test collect(gy) ≈ m(x)
+end
diff --git a/test/layers/upsample.jl b/test/layers/upsample.jl
index 936767dcc4..3605a3f597 100644
--- a/test/layers/upsample.jl
+++ b/test/layers/upsample.jl
@@ -13,6 +13,12 @@
 end
 
 @testset "upsample nearest" begin
+  x = rand(Float32, 3, 2, 3)
+  m = Upsample(2, mode=:nearest)
+  y = m(x)
+  @test y isa Array{Float32, 3} 
+  @test size(y) == (6, 2, 3)
+
   x = rand(Float32, 3, 4, 2, 3)
   m = Upsample(scale=(2, 3), mode=:nearest)
   y = m(x)
@@ -24,12 +30,6 @@ end
   y = m(x)
   @test y isa Array{Float32, 4} 
   @test size(y) == (6, 8, 2, 3)
-
-  x = rand(Float32, 3, 2, 3)
-  m = Upsample(2, mode=:nearest)
-  y = m(x)
-  @test y isa Array{Float32, 3} 
-  @test size(y) == (6, 2, 3)
 end
 
 @testset "PixelShuffle" begin

From f531206321569b75086b2076b32f938c2bad54d4 Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Thu, 21 Jan 2021 02:46:30 +0100
Subject: [PATCH 08/19] arg error

---
 src/layers/upsample.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/layers/upsample.jl b/src/layers/upsample.jl
index 52b9b660f1..cda4859649 100644
--- a/src/layers/upsample.jl
+++ b/src/layers/upsample.jl
@@ -31,7 +31,7 @@ Upsample(; scale, mode::Symbol=:nearest) = Upsample(scale; mode)
 
 function Upsample(scale; mode::Symbol=:nearest)
   mode in [:nearest, :bilinear] || 
-    throw(ArgumentError("mode=:$mode is not supported, allowed values are :nearest, :bilinear"))
+    throw(ArgumentError("mode=:$mode is not supported."))
   return Upsample{mode, typeof(scale)}(scale)
 end
 

From 93287ddd8655ff8b88aded53f3d2962b4e70720f Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Fri, 5 Feb 2021 09:36:26 +0100
Subject: [PATCH 09/19] update

---
 src/layers/upsample.jl | 47 ++++++++++++++++++++++++++----------------
 1 file changed, 29 insertions(+), 18 deletions(-)

diff --git a/src/layers/upsample.jl b/src/layers/upsample.jl
index cda4859649..fcd0521b21 100644
--- a/src/layers/upsample.jl
+++ b/src/layers/upsample.jl
@@ -1,20 +1,25 @@
 """
   Upsample(scale; mode=:nearest)  
-  Upsample(; scale, mode=:nearest)
+  Upsample(; size, mode=:bilinear)
 
 An upsampling layer. `scale` is a number or a tuple of numbers 
 representing  the output rescaling factor along each spatial dimension.
+
 Currently supported upsampling `mode`s are:
   - `:nearest`
   - `:bilinear`
 
+For some `mode`s it is possible
+to directly specify the output spatial `size`,
+as an alternative to `scale`,   
+
 See [`NNlib.upsample_nearest`](@ref), [`NNlib.upsample_bilinear`](@ref).
 
 # Examples
 
 ```juliarepl
-julia> m = Upsample(scale=(2,3), mode=:bilinear)
-Upsample(scale=(2, 3), mode=bilinear)
+julia> m = Upsample((2,3), mode=:bilinear)
+Upsample((2, 3), mode=bilinear)
 
 julia> m(ones(1,1,1,1))
 2×3×1×1 Array{Float64,4}:
@@ -23,35 +28,39 @@ julia> m(ones(1,1,1,1))
  1.0  1.0  1.0
 ```
 """
-struct Upsample{Mode,T}
-  scale::T
+struct Upsample{Mode,S,T}
+  scale::S
+  size::T
 end
 
-Upsample(; scale, mode::Symbol=:nearest) = Upsample(scale; mode)
-
 function Upsample(scale; mode::Symbol=:nearest)
   mode in [:nearest, :bilinear] || 
     throw(ArgumentError("mode=:$mode is not supported."))
-  return Upsample{mode, typeof(scale)}(scale)
+  return Upsample(scale, nothing)
+end
+
+function Upsample(; size, mode::Symbol=:bilinear)
+  mode in [:nearest, :bilinear] || 
+    throw(ArgumentError("mode=:$mode is not supported."))
+  return Upsample(nothing, size)
 end
 
-(m::Upsample{:nearest,<:Number})(x::AbstractArray{T,N}) where {T,N} =
-  NNlib.upsample_nearest(x, ntuple(_ -> m.scale, N-2))
-(m::Upsample{:nearest,<:Tuple})(x::AbstractArray) =
+(m::Upsample{:nearest})(x::AbstractArray) =
   NNlib.upsample_nearest(x, m.scale)
 
-(m::Upsample{:bilinear,<:Number})(x::AbstractArray{T,N}) where {T,N} =
-  NNlib.upsample_bilinear(x, ntuple(_ -> m.scale, N-2))
-(m::Upsample{:bilinear,<:Tuple})(x::AbstractArray) =
+(m::Upsample{:bilinear})(x::AbstractArray) =
   NNlib.upsample_bilinear(x, m.scale)
+(m::Upsample{:bilinear, Nothing})(x::AbstractArray) = 
+  NNlib.upsample_bilinear(x; size=m.size)
 
 function Base.show(io::IO, u::Upsample{mode}) where {mode}
-  print(io, "Upsample(scale=", u.scale)
-  print(io, ", mode=:", mode)
-  print(io, ")")
+  print(io, "Upsample(")
+  u.scale !== nothing && print(io, "$(u.scale), ")
+  u.size !== nothing && print(io, "size=$(u.size), ")
+  print(io, "mode=:", mode)
+  println(io, ")")
 end
 
-
 """
     PixelShuffle(r::Int)
 
@@ -64,3 +73,5 @@ struct PixelShuffle
 end
 
 (m::PixelShuffle)(x) = NNlib.pixel_shuffle(x, m.r)
+
+

From bf06f0c25386004232814d205b80288318f1fedd Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Sun, 7 Feb 2021 09:27:41 +0100
Subject: [PATCH 10/19] update to latest nnlib

---
 Manifest.toml           |  4 ++--
 Project.toml            |  2 +-
 src/layers/upsample.jl  | 37 ++++++++++++++++++++++---------------
 test/cuda/layers.jl     | 32 +++++++++-----------------------
 test/layers/upsample.jl | 13 +++++++++----
 5 files changed, 43 insertions(+), 45 deletions(-)

diff --git a/Manifest.toml b/Manifest.toml
index e3b463a9a2..7210fa33bf 100644
--- a/Manifest.toml
+++ b/Manifest.toml
@@ -236,9 +236,9 @@ uuid = "a63ad114-7e13-5084-954f-fe012c677804"
 
 [[NNlib]]
 deps = ["ChainRulesCore", "Compat", "LinearAlgebra", "Pkg", "Requires", "Statistics"]
-git-tree-sha1 = "e031936930ecb559a92992706f8acb02dc81fed2"
+git-tree-sha1 = "df42d0816edfc24f5b82a728f46381613c4dff79"
 uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
-version = "0.7.13"
+version = "0.7.14"
 
 [[NaNMath]]
 git-tree-sha1 = "bfe47e760d60b82b66b61d2d44128b62e3a369fb"
diff --git a/Project.toml b/Project.toml
index 17ffe082bb..b51f143252 100644
--- a/Project.toml
+++ b/Project.toml
@@ -34,7 +34,7 @@ Colors = "0.12"
 Functors = "0.1, 0.2"
 Juno = "0.8"
 MacroTools = "0.5"
-NNlib = "0.7.10"
+NNlib = "0.7.14"
 Reexport = "0.2, 1.0"
 StatsBase = "0.33"
 ZipFile = "0.9"
diff --git a/src/layers/upsample.jl b/src/layers/upsample.jl
index fcd0521b21..b2a799a346 100644
--- a/src/layers/upsample.jl
+++ b/src/layers/upsample.jl
@@ -1,25 +1,27 @@
 """
   Upsample(scale; mode=:nearest)  
-  Upsample(; size, mode=:bilinear)
+  Upsample(; size, mode=:nearest)
 
-An upsampling layer. `scale` is a number or a tuple of numbers 
-representing  the output rescaling factor along each spatial dimension.
+An upsampling layer. 
 
-Currently supported upsampling `mode`s are:
-  - `:nearest`
-  - `:bilinear`
+`scale` is a number or a tuple of numbers 
+representing  the output rescaling factor along each spatial dimension.
+For integer `scale`, all but the last 2 dimensions (channel and batch)
+will rescaled by the same factor. 
 
-For some `mode`s it is possible
-to directly specify the output spatial `size`,
-as an alternative to `scale`,   
+It is also possible to directly specify the output spatial `size`,
+as an alternative to using `scale`.
 
-See [`NNlib.upsample_nearest`](@ref), [`NNlib.upsample_bilinear`](@ref).
+Currently supported upsampling `mode`s 
+and corresponding NNlib's methods are:
+  - `:nearest` -> [`NNlib.upsample_nearest`](@ref) 
+  - `:bilinear` -> [`NNlib.upsample_bilinear`](@ref)
 
 # Examples
 
 ```juliarepl
 julia> m = Upsample((2,3), mode=:bilinear)
-Upsample((2, 3), mode=bilinear)
+Upsample((2, 3), mode=:bilinear)
 
 julia> m(ones(1,1,1,1))
 2×3×1×1 Array{Float64,4}:
@@ -33,20 +35,25 @@ struct Upsample{Mode,S,T}
   size::T
 end
 
-function Upsample(scale; mode::Symbol=:nearest)
+function Upsample(scale::S; mode::Symbol=:nearest) where S
   mode in [:nearest, :bilinear] || 
     throw(ArgumentError("mode=:$mode is not supported."))
-  return Upsample(scale, nothing)
+  return Upsample{mode,S,Nothing}(scale, nothing)
 end
 
-function Upsample(; size, mode::Symbol=:bilinear)
+function Upsample(; size::T, mode::Symbol=:nearest) where T
   mode in [:nearest, :bilinear] || 
     throw(ArgumentError("mode=:$mode is not supported."))
-  return Upsample(nothing, size)
+  return Upsample{mode, Nothing, T}(nothing, size)
 end
 
 (m::Upsample{:nearest})(x::AbstractArray) =
   NNlib.upsample_nearest(x, m.scale)
+function (m::Upsample{:nearest, Int})(x::AbstractArray{T, N}) where {T, N} 
+  NNlib.upsample_nearest(x, ntuple(i -> m.scale, N-2))
+end
+(m::Upsample{:nearest, Nothing})(x::AbstractArray) =
+  NNlib.upsample_nearest(x; size=m.size)
 
 (m::Upsample{:bilinear})(x::AbstractArray) =
   NNlib.upsample_bilinear(x, m.scale)
diff --git a/test/cuda/layers.jl b/test/cuda/layers.jl
index 3c44f5cab3..f615ede265 100644
--- a/test/cuda/layers.jl
+++ b/test/cuda/layers.jl
@@ -79,6 +79,15 @@ gpu_gradtest("GroupNorm 3d", groupnorm, rand(Float32, 8, 8, 8, 12, 4), 12, 3, se
 gpu_gradtest("GroupNorm 2d", groupnorm, rand(Float32, 8, 8, 12, 4), 12, 3, setmode=true)
 gpu_gradtest("GroupNorm 1d", groupnorm, rand(Float32, 8, 3, 12, 4), 12, 3, setmode=true)
 
+upsample = [Upsample]
+gpu_gradtest("Upsample 2d", upsample, rand(Float32, 3, 4, 2, 3), (2,2))
+gpu_gradtest("Upsample 1d", upsample, rand(Float32, 3, 4, 2, 3), (2,))
+
+pixelshuffle = [PixelShuffle]
+gpu_gradtest("PixelShuffle 2d", pixelshuffle, rand(Float32, 3, 4, 18, 3), 3)
+gpu_gradtest("PixelShuffle 1d", pixelshuffle, rand(Float32, 3, 18, 3), 3)
+
+
 @testset "function layers" begin
   x = rand(Float32, 3,3)
   gpu_autodiff_test(x -> sum(Flux.normalise(x; dims=1)), x)
@@ -169,26 +178,3 @@ end
   gs = gradient(() -> sum(l(ip)), Flux.params(l))
   @test l.b ∉ gs.params 
 end
-
-@testset "upsample bilinear" begin
-  m = Upsample(3, mode=:bilinear)
-  x = rand(Float32, 3, 4, 2, 3)
-  gy = m(x |> gpu)
-  @test gy isa CuArray{Float32, 4}
-  @test collect(gy) ≈ m(x)
-end
-
-@testset "upsample nearest" begin
-  x = rand(Float32, 3, 4, 2, 3)
-  gy = m(x |> gpu)
-  @test gy isa CuArray{Float32, 4}
-  @test collect(gy) ≈ m(x)
-end
-
-@testset "PixelShuffle" begin
-  m = PixelShuffle(3)
-  x = rand(Float32, 3, 4, 18, 3)
-  gy = m(x |> gpu)
-  @test gy isa CuArray{Float32, 4}
-  @test collect(gy) ≈ m(x)
-end
diff --git a/test/layers/upsample.jl b/test/layers/upsample.jl
index 3605a3f597..40702e0930 100644
--- a/test/layers/upsample.jl
+++ b/test/layers/upsample.jl
@@ -1,5 +1,5 @@
 @testset "upsample bilinear" begin
-  m = Upsample(scale=(2, 3), mode=:bilinear)
+  m = Upsample((2, 3), mode=:bilinear)
   x = rand(Float32, 3, 4, 2, 3)
   y = m(x)
   @test y isa Array{Float32, 4} 
@@ -14,18 +14,23 @@ end
 
 @testset "upsample nearest" begin
   x = rand(Float32, 3, 2, 3)
-  m = Upsample(2, mode=:nearest)
+  m = Upsample((2,), mode=:nearest)
   y = m(x)
   @test y isa Array{Float32, 3} 
   @test size(y) == (6, 2, 3)
 
   x = rand(Float32, 3, 4, 2, 3)
-  m = Upsample(scale=(2, 3), mode=:nearest)
+  
+  m = Upsample((2, 3), mode=:nearest)
   y = m(x)
   @test y isa Array{Float32, 4} 
   @test size(y) == (6, 12, 2, 3)
   
-  x = rand(Float32, 3, 4, 2, 3)
+  m = Upsample((2,), mode=:nearest)
+  y = m(x)
+  @test y isa Array{Float32, 4} 
+  @test size(y) == (6, 4, 2, 3)
+
   m = Upsample(2, mode=:nearest)
   y = m(x)
   @test y isa Array{Float32, 4} 

From 6d1d55f58aa17c19faada694ea16b8ad7422ce91 Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Sun, 7 Feb 2021 10:29:59 +0100
Subject: [PATCH 11/19] fix outputsize

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

diff --git a/src/outputsize.jl b/src/outputsize.jl
index eaee34433f..7502627929 100644
--- a/src/outputsize.jl
+++ b/src/outputsize.jl
@@ -35,7 +35,7 @@ Base.isless(::Nil, ::Number) = true
 Base.isless(::Number, ::Nil) = true
 
 Base.isnan(::Nil) = false
-
+Base.isfinite(::Nil) = true
 Base.typemin(::Type{Nil}) = nil
 Base.typemax(::Type{Nil}) = nil
 
@@ -100,7 +100,7 @@ nil_input(pad::Bool, tup::Tuple{Vararg{Tuple}}) = nil_input(pad, tup...)
 function outputsize(m::Chain, inputsizes::Tuple{Vararg{Integer}}...; padbatch=false)
   x = nil_input(padbatch, inputsizes...)
   for (i,lay) in enumerate(m.layers)
-    try
+    try 
       x = lay(x)
     catch err
       str = x isa AbstractArray ? "with input of size $(size(x))" : ""
diff --git a/test/outputsize.jl b/test/outputsize.jl
index 5b25c8bbd8..eec96023df 100644
--- a/test/outputsize.jl
+++ b/test/outputsize.jl
@@ -146,9 +146,7 @@ end
   @test outputsize(m, (32, 32, 3, 16)) == (32, 32, 3, 16)
   @test outputsize(m, (32, 32, 3); padbatch=true) == (32, 32, 3, 1)
 
-  if VERSION >= v"1.1"
-    m = GroupNorm(16, 4)
-    @test outputsize(m, (32, 32, 16, 16)) == (32, 32, 16, 16)
-    @test outputsize(m, (32, 32, 16); padbatch=true) == (32, 32, 16, 1)
-  end
+  m = GroupNorm(16, 4)
+  @test outputsize(m, (32, 32, 16, 16)) == (32, 32, 16, 16)
+  @test outputsize(m, (32, 32, 16); padbatch=true) == (32, 32, 16, 1)
 end

From e252e460ca009fa58583a853c40e2b78a0676634 Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Wed, 10 Feb 2021 09:46:45 +0100
Subject: [PATCH 12/19] new interface

---
 NEWS.md                 |  1 -
 src/layers/upsample.jl  | 42 ++++++++++++++++++++---------------------
 src/outputsize.jl       |  2 +-
 test/cuda/layers.jl     |  2 +-
 test/layers/upsample.jl | 23 ++++++++++++++++------
 5 files changed, 39 insertions(+), 31 deletions(-)

diff --git a/NEWS.md b/NEWS.md
index 6f77ede84f..ceb10bdcee 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -13,7 +13,6 @@
 * New [`Parallel` layer](https://github.com/FluxML/Flux.jl/pull/1462) adds inception module-like building blocks.
 * Feature additions and bug fixes for BatchNorm, LayerNorm, InstanceNorm, and GroupNorm [normalization layers](https://github.com/FluxML/Flux.jl/pull/1397)
 * Added [Upsample and PixelShuffle layers](https://github.com/FluxML/Flux.jl/pull/1468)
-* Other new features and bug fixes (see GitHub releases page)
 
 ## v0.11.2
 
diff --git a/src/layers/upsample.jl b/src/layers/upsample.jl
index b2a799a346..f462fa16f1 100644
--- a/src/layers/upsample.jl
+++ b/src/layers/upsample.jl
@@ -1,6 +1,5 @@
 """
-  Upsample(scale; mode=:nearest)  
-  Upsample(; size, mode=:nearest)
+  Upsample(mode=:nearest; scale=nothing, size=nothing)  
 
 An upsampling layer. 
 
@@ -20,31 +19,30 @@ and corresponding NNlib's methods are:
 # Examples
 
 ```juliarepl
-julia> m = Upsample((2,3), mode=:bilinear)
-Upsample((2, 3), mode=:bilinear)
-
-julia> m(ones(1,1,1,1))
-2×3×1×1 Array{Float64,4}:
-[:, :, 1, 1] =
- 1.0  1.0  1.0
- 1.0  1.0  1.0
-```
+julia> m = Upsample(scale=(2, 3))
+Upsample(:nearest, scale=(2, 3))
+
+julia> m(ones(2, 2, 1, 1)) |> size
+(4, 6, 1, 1)
+
+julia> m = Upsample(:bilinear, size=(4, 5))
+Upsample(:bilinear, size=(4, 5))
+
+julia> m(ones(2, 2, 1, 1)) |> size
+(4, 5, 1, 1)
 """
 struct Upsample{Mode,S,T}
   scale::S
   size::T
 end
 
-function Upsample(scale::S; mode::Symbol=:nearest) where S
-  mode in [:nearest, :bilinear] || 
-    throw(ArgumentError("mode=:$mode is not supported."))
-  return Upsample{mode,S,Nothing}(scale, nothing)
-end
-
-function Upsample(; size::T, mode::Symbol=:nearest) where T
+function Upsample(mode::Symbol=:nearest; scale=nothing, size=nothing)
   mode in [:nearest, :bilinear] || 
     throw(ArgumentError("mode=:$mode is not supported."))
-  return Upsample{mode, Nothing, T}(nothing, size)
+  if ~((scale === nothing) ⊻ (size === nothing))
+    throw(ArgumentError("Either scale or size should be specified."))
+  end
+  return Upsample{mode,typeof(scale),typeof(size)}(scale, size)
 end
 
 (m::Upsample{:nearest})(x::AbstractArray) =
@@ -62,9 +60,9 @@ end
 
 function Base.show(io::IO, u::Upsample{mode}) where {mode}
   print(io, "Upsample(")
-  u.scale !== nothing && print(io, "$(u.scale), ")
-  u.size !== nothing && print(io, "size=$(u.size), ")
-  print(io, "mode=:", mode)
+  print(io, ":", mode)
+  u.scale !== nothing && print(io, ", scale=$(u.scale)")
+  u.size !== nothing && print(io, ", size=$(u.size)")
   println(io, ")")
 end
 
diff --git a/src/outputsize.jl b/src/outputsize.jl
index 7502627929..303e2c9c8b 100644
--- a/src/outputsize.jl
+++ b/src/outputsize.jl
@@ -100,7 +100,7 @@ nil_input(pad::Bool, tup::Tuple{Vararg{Tuple}}) = nil_input(pad, tup...)
 function outputsize(m::Chain, inputsizes::Tuple{Vararg{Integer}}...; padbatch=false)
   x = nil_input(padbatch, inputsizes...)
   for (i,lay) in enumerate(m.layers)
-    try 
+    try
       x = lay(x)
     catch err
       str = x isa AbstractArray ? "with input of size $(size(x))" : ""
diff --git a/test/cuda/layers.jl b/test/cuda/layers.jl
index f615ede265..6a7dcd1b05 100644
--- a/test/cuda/layers.jl
+++ b/test/cuda/layers.jl
@@ -79,7 +79,7 @@ gpu_gradtest("GroupNorm 3d", groupnorm, rand(Float32, 8, 8, 8, 12, 4), 12, 3, se
 gpu_gradtest("GroupNorm 2d", groupnorm, rand(Float32, 8, 8, 12, 4), 12, 3, setmode=true)
 gpu_gradtest("GroupNorm 1d", groupnorm, rand(Float32, 8, 3, 12, 4), 12, 3, setmode=true)
 
-upsample = [Upsample]
+upsample = [x -> Upsample(scale=x)]
 gpu_gradtest("Upsample 2d", upsample, rand(Float32, 3, 4, 2, 3), (2,2))
 gpu_gradtest("Upsample 1d", upsample, rand(Float32, 3, 4, 2, 3), (2,))
 
diff --git a/test/layers/upsample.jl b/test/layers/upsample.jl
index 40702e0930..1419f44409 100644
--- a/test/layers/upsample.jl
+++ b/test/layers/upsample.jl
@@ -1,37 +1,48 @@
 @testset "upsample bilinear" begin
-  m = Upsample((2, 3), mode=:bilinear)
+  m = Upsample(:bilinear, scale=(2, 3))
   x = rand(Float32, 3, 4, 2, 3)
   y = m(x)
   @test y isa Array{Float32, 4} 
   @test size(y) == (6, 12, 2, 3)
 
-  m = Upsample(3, mode=:bilinear)
+  m = Upsample(:bilinear, scale=3)
   x = rand(Float32, 3, 4, 2, 3)
   y = m(x)
   @test y isa Array{Float32, 4} 
   @test size(y) == (9, 12, 2, 3)
+
+  m = Upsample(:bilinear, size=(4, 6))
+  x = rand(Float32, 3, 4, 2, 3)
+  y = m(x)
+  @test y isa Array{Float32, 4} 
+  @test size(y) == (4, 6, 2, 3)
 end
 
 @testset "upsample nearest" begin
   x = rand(Float32, 3, 2, 3)
-  m = Upsample((2,), mode=:nearest)
+  m = Upsample(:nearest, scale=(2,))
   y = m(x)
   @test y isa Array{Float32, 3} 
   @test size(y) == (6, 2, 3)
 
   x = rand(Float32, 3, 4, 2, 3)
   
-  m = Upsample((2, 3), mode=:nearest)
+  m = Upsample(:nearest, scale=(2, 3))
   y = m(x)
   @test y isa Array{Float32, 4} 
   @test size(y) == (6, 12, 2, 3)
   
-  m = Upsample((2,), mode=:nearest)
+  m = Upsample(:nearest, scale=(2,))
   y = m(x)
   @test y isa Array{Float32, 4} 
   @test size(y) == (6, 4, 2, 3)
 
-  m = Upsample(2, mode=:nearest)
+  m = Upsample(:nearest, scale=2)
+  y = m(x)
+  @test y isa Array{Float32, 4} 
+  @test size(y) == (6, 8, 2, 3)
+
+  m = Upsample(:nearest, size=(6,8))
   y = m(x)
   @test y isa Array{Float32, 4} 
   @test size(y) == (6, 8, 2, 3)

From 9c9149163f398e0b0e99b820f302ebdf5194a307 Mon Sep 17 00:00:00 2001
From: Dhairya Gandhi <dhairya@juliacomputing.com>
Date: Wed, 10 Feb 2021 16:40:15 +0530
Subject: [PATCH 13/19] style cleanups

---
 src/layers/upsample.jl | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/src/layers/upsample.jl b/src/layers/upsample.jl
index f462fa16f1..0a68a1f956 100644
--- a/src/layers/upsample.jl
+++ b/src/layers/upsample.jl
@@ -1,12 +1,12 @@
 """
-  Upsample(mode=:nearest; scale=nothing, size=nothing)  
+  Upsample(mode = :nearest; scale = nothing, size = nothing)  
 
 An upsampling layer. 
 
 `scale` is a number or a tuple of numbers 
 representing  the output rescaling factor along each spatial dimension.
 For integer `scale`, all but the last 2 dimensions (channel and batch)
-will rescaled by the same factor. 
+will be rescaled by the same factor. 
 
 It is also possible to directly specify the output spatial `size`,
 as an alternative to using `scale`.
@@ -19,13 +19,13 @@ and corresponding NNlib's methods are:
 # Examples
 
 ```juliarepl
-julia> m = Upsample(scale=(2, 3))
+julia> m = Upsample(scale = (2, 3))
 Upsample(:nearest, scale=(2, 3))
 
 julia> m(ones(2, 2, 1, 1)) |> size
 (4, 6, 1, 1)
 
-julia> m = Upsample(:bilinear, size=(4, 5))
+julia> m = Upsample(:bilinear, size = (4, 5))
 Upsample(:bilinear, size=(4, 5))
 
 julia> m(ones(2, 2, 1, 1)) |> size
@@ -36,7 +36,7 @@ struct Upsample{Mode,S,T}
   size::T
 end
 
-function Upsample(mode::Symbol=:nearest; scale=nothing, size=nothing)
+function Upsample(mode = :nearest; scale = nothing, size = nothing)
   mode in [:nearest, :bilinear] || 
     throw(ArgumentError("mode=:$mode is not supported."))
   if ~((scale === nothing) ⊻ (size === nothing))
@@ -61,8 +61,8 @@ end
 function Base.show(io::IO, u::Upsample{mode}) where {mode}
   print(io, "Upsample(")
   print(io, ":", mode)
-  u.scale !== nothing && print(io, ", scale=$(u.scale)")
-  u.size !== nothing && print(io, ", size=$(u.size)")
+  u.scale !== nothing && print(io, ", scale = $(u.scale)")
+  u.size !== nothing && print(io, ", size = $(u.size)")
   println(io, ")")
 end
 
@@ -79,4 +79,3 @@ end
 
 (m::PixelShuffle)(x) = NNlib.pixel_shuffle(x, m.r)
 
-

From e2bc4e8638a1dd9b9ac264eb65a7e3b128b347e6 Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@unibocconi.it>
Date: Wed, 10 Feb 2021 22:59:44 +0100
Subject: [PATCH 14/19] update docstring

Co-authored-by: Michael Abbott <32575566+mcabbott@users.noreply.github.com>
---
 src/layers/upsample.jl | 15 +++++----------
 1 file changed, 5 insertions(+), 10 deletions(-)

diff --git a/src/layers/upsample.jl b/src/layers/upsample.jl
index 0a68a1f956..b1d322f2ed 100644
--- a/src/layers/upsample.jl
+++ b/src/layers/upsample.jl
@@ -1,15 +1,11 @@
 """
-  Upsample(mode = :nearest; scale = nothing, size = nothing)  
+  Upsample(mode = :nearest; [scale, size])  
 
-An upsampling layer. 
+An upsampling layer. One of two keywords must be given:
 
-`scale` is a number or a tuple of numbers 
-representing  the output rescaling factor along each spatial dimension.
-For integer `scale`, all but the last 2 dimensions (channel and batch)
-will be rescaled by the same factor. 
-
-It is also possible to directly specify the output spatial `size`,
-as an alternative to using `scale`.
+If `scale` is a number, this applies to all but the last two dimensions (channel and batch) of the input. 
+It may also be a tuple, to control dimensions individually. Alternatively, keyword 
+`size` accepts a tuple, to directly specify the leading dimensions of the output.
 
 Currently supported upsampling `mode`s 
 and corresponding NNlib's methods are:
@@ -78,4 +74,3 @@ struct PixelShuffle
 end
 
 (m::PixelShuffle)(x) = NNlib.pixel_shuffle(x, m.r)
-

From 05a9a2d810a86dbf16916dd66583d0e9eb74cd8b Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Wed, 10 Feb 2021 23:36:29 +0100
Subject: [PATCH 15/19] add other constructor + small changes

---
 src/layers/upsample.jl | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/src/layers/upsample.jl b/src/layers/upsample.jl
index b1d322f2ed..5b76a742b7 100644
--- a/src/layers/upsample.jl
+++ b/src/layers/upsample.jl
@@ -1,5 +1,6 @@
 """
-  Upsample(mode = :nearest; [scale, size])  
+  Upsample(mode = :nearest; [scale, size]) 
+  Upsample(scale, mode = :nearest)  
 
 An upsampling layer. One of two keywords must be given:
 
@@ -16,31 +17,33 @@ and corresponding NNlib's methods are:
 
 ```juliarepl
 julia> m = Upsample(scale = (2, 3))
-Upsample(:nearest, scale=(2, 3))
+Upsample(:nearest, scale = (2, 3))
 
 julia> m(ones(2, 2, 1, 1)) |> size
 (4, 6, 1, 1)
 
 julia> m = Upsample(:bilinear, size = (4, 5))
-Upsample(:bilinear, size=(4, 5))
+Upsample(:bilinear, size = (4, 5))
 
 julia> m(ones(2, 2, 1, 1)) |> size
 (4, 5, 1, 1)
 """
-struct Upsample{Mode,S,T}
+struct Upsample{mode, S, T}
   scale::S
   size::T
 end
 
-function Upsample(mode = :nearest; scale = nothing, size = nothing)
+function Upsample(mode::Symbol = :nearest; scale = nothing, size = nothing)
   mode in [:nearest, :bilinear] || 
     throw(ArgumentError("mode=:$mode is not supported."))
-  if ~((scale === nothing) ⊻ (size === nothing))
-    throw(ArgumentError("Either scale or size should be specified."))
+  if !(isnothing(scale) ⊻ isnothing(size))
+      throw(ArgumentError("Either scale or size should be specified (but not both)."))
   end
   return Upsample{mode,typeof(scale),typeof(size)}(scale, size)
 end
 
+Upsample(scale, mode::Symbol = :nearest) = Upsample(mode; scale)
+
 (m::Upsample{:nearest})(x::AbstractArray) =
   NNlib.upsample_nearest(x, m.scale)
 function (m::Upsample{:nearest, Int})(x::AbstractArray{T, N}) where {T, N} 

From 8503f4be1404962807bf816d745395dc86a8e9c4 Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Wed, 10 Feb 2021 23:38:24 +0100
Subject: [PATCH 16/19] add test

---
 test/layers/upsample.jl | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/test/layers/upsample.jl b/test/layers/upsample.jl
index 1419f44409..0d5221a316 100644
--- a/test/layers/upsample.jl
+++ b/test/layers/upsample.jl
@@ -42,6 +42,9 @@ end
   @test y isa Array{Float32, 4} 
   @test size(y) == (6, 8, 2, 3)
 
+  m = Upsample(2)
+  @test y2 ≈ y 
+
   m = Upsample(:nearest, size=(6,8))
   y = m(x)
   @test y isa Array{Float32, 4} 

From 91cac26e226c6b12104fb1cf2098b1abaeab38d1 Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Wed, 10 Feb 2021 23:45:37 +0100
Subject: [PATCH 17/19] fix test

---
 test/layers/upsample.jl |  1 +
 test/runtests.jl        | 84 ++++++++++++++++++++---------------------
 2 files changed, 43 insertions(+), 42 deletions(-)

diff --git a/test/layers/upsample.jl b/test/layers/upsample.jl
index 0d5221a316..6b6b6bc032 100644
--- a/test/layers/upsample.jl
+++ b/test/layers/upsample.jl
@@ -43,6 +43,7 @@ end
   @test size(y) == (6, 8, 2, 3)
 
   m = Upsample(2)
+  y2 = m(x)
   @test y2 ≈ y 
 
   m = Upsample(:nearest, size=(6,8))
diff --git a/test/runtests.jl b/test/runtests.jl
index a40433d0f1..86579d2bef 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -6,49 +6,49 @@ using IterTools: ncycle
 
 Random.seed!(0)
 
-@testset "Utils" begin
-  include("utils.jl")
-end
-
-@testset "Onehot" begin
-  include("onehot.jl")
-end
-
-@testset "Optimise" begin
-  include("optimise.jl")
-end
-
-@testset "Data" begin
-  include("data.jl")
-end
-
-@testset "Losses" begin
-  include("losses.jl")
-  include("ctc.jl")
-  if Flux.use_cuda[] include("ctc-gpu.jl") end
-end
-
-@testset "Layers" begin
-  include("layers/basic.jl")
-  include("layers/normalisation.jl")
-  include("layers/stateless.jl")
-  include("layers/recurrent.jl")
-  include("layers/conv.jl")
+# @testset "Utils" begin
+#   include("utils.jl")
+# end
+
+# @testset "Onehot" begin
+#   include("onehot.jl")
+# end
+
+# @testset "Optimise" begin
+#   include("optimise.jl")
+# end
+
+# @testset "Data" begin
+#   include("data.jl")
+# end
+
+# @testset "Losses" begin
+#   include("losses.jl")
+#   include("ctc.jl")
+#   if Flux.use_cuda[] include("ctc-gpu.jl") end
+# end
+
+# @testset "Layers" begin
+#   include("layers/basic.jl")
+#   include("layers/normalisation.jl")
+#   include("layers/stateless.jl")
+#   include("layers/recurrent.jl")
+#   include("layers/conv.jl")
   include("layers/upsample.jl")
-end
-
-@testset "outputsize" begin
-  using Flux: outputsize
-  include("outputsize.jl")
-end
-
-@testset "CUDA" begin
-  if Flux.use_cuda[]
-    include("cuda/runtests.jl")
-  else
-    @warn "CUDA unavailable, not testing GPU support"
-  end
-end
+# end
+
+# @testset "outputsize" begin
+#   using Flux: outputsize
+#   include("outputsize.jl")
+# end
+
+# @testset "CUDA" begin
+#   if Flux.use_cuda[]
+#     include("cuda/runtests.jl")
+#   else
+#     @warn "CUDA unavailable, not testing GPU support"
+#   end
+# end
 
 @static if VERSION == v"1.5"
   using Documenter

From 75094d6d4b2368c3e7b75513e4f617190566f8c2 Mon Sep 17 00:00:00 2001
From: Carlo Lucibello <carlo.lucibello@gmail.com>
Date: Wed, 10 Feb 2021 23:46:24 +0100
Subject: [PATCH 18/19] fix test

---
 src/layers/upsample.jl |  2 +-
 test/runtests.jl       | 84 +++++++++++++++++++++---------------------
 2 files changed, 43 insertions(+), 43 deletions(-)

diff --git a/src/layers/upsample.jl b/src/layers/upsample.jl
index 5b76a742b7..f2134f2168 100644
--- a/src/layers/upsample.jl
+++ b/src/layers/upsample.jl
@@ -37,7 +37,7 @@ function Upsample(mode::Symbol = :nearest; scale = nothing, size = nothing)
   mode in [:nearest, :bilinear] || 
     throw(ArgumentError("mode=:$mode is not supported."))
   if !(isnothing(scale) ⊻ isnothing(size))
-      throw(ArgumentError("Either scale or size should be specified (but not both)."))
+    throw(ArgumentError("Either scale or size should be specified (but not both)."))
   end
   return Upsample{mode,typeof(scale),typeof(size)}(scale, size)
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index 86579d2bef..a40433d0f1 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -6,49 +6,49 @@ using IterTools: ncycle
 
 Random.seed!(0)
 
-# @testset "Utils" begin
-#   include("utils.jl")
-# end
-
-# @testset "Onehot" begin
-#   include("onehot.jl")
-# end
-
-# @testset "Optimise" begin
-#   include("optimise.jl")
-# end
-
-# @testset "Data" begin
-#   include("data.jl")
-# end
-
-# @testset "Losses" begin
-#   include("losses.jl")
-#   include("ctc.jl")
-#   if Flux.use_cuda[] include("ctc-gpu.jl") end
-# end
-
-# @testset "Layers" begin
-#   include("layers/basic.jl")
-#   include("layers/normalisation.jl")
-#   include("layers/stateless.jl")
-#   include("layers/recurrent.jl")
-#   include("layers/conv.jl")
+@testset "Utils" begin
+  include("utils.jl")
+end
+
+@testset "Onehot" begin
+  include("onehot.jl")
+end
+
+@testset "Optimise" begin
+  include("optimise.jl")
+end
+
+@testset "Data" begin
+  include("data.jl")
+end
+
+@testset "Losses" begin
+  include("losses.jl")
+  include("ctc.jl")
+  if Flux.use_cuda[] include("ctc-gpu.jl") end
+end
+
+@testset "Layers" begin
+  include("layers/basic.jl")
+  include("layers/normalisation.jl")
+  include("layers/stateless.jl")
+  include("layers/recurrent.jl")
+  include("layers/conv.jl")
   include("layers/upsample.jl")
-# end
-
-# @testset "outputsize" begin
-#   using Flux: outputsize
-#   include("outputsize.jl")
-# end
-
-# @testset "CUDA" begin
-#   if Flux.use_cuda[]
-#     include("cuda/runtests.jl")
-#   else
-#     @warn "CUDA unavailable, not testing GPU support"
-#   end
-# end
+end
+
+@testset "outputsize" begin
+  using Flux: outputsize
+  include("outputsize.jl")
+end
+
+@testset "CUDA" begin
+  if Flux.use_cuda[]
+    include("cuda/runtests.jl")
+  else
+    @warn "CUDA unavailable, not testing GPU support"
+  end
+end
 
 @static if VERSION == v"1.5"
   using Documenter

From 9acaae9dbdf1f0b69bebd105cbb25e70b80b88cb Mon Sep 17 00:00:00 2001
From: Dhairya Gandhi <dhairya@juliacomputing.com>
Date: Thu, 11 Feb 2021 14:02:21 +0530
Subject: [PATCH 19/19] Println -> print

Co-authored-by: Michael Abbott <32575566+mcabbott@users.noreply.github.com>
---
 src/layers/upsample.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/layers/upsample.jl b/src/layers/upsample.jl
index f2134f2168..51d3fdb277 100644
--- a/src/layers/upsample.jl
+++ b/src/layers/upsample.jl
@@ -62,7 +62,7 @@ function Base.show(io::IO, u::Upsample{mode}) where {mode}
   print(io, ":", mode)
   u.scale !== nothing && print(io, ", scale = $(u.scale)")
   u.size !== nothing && print(io, ", size = $(u.size)")
-  println(io, ")")
+  print(io, ")")
 end
 
 """