From 9cad84bc9ec4699a76c002e15ce985a5dc61f9bf Mon Sep 17 00:00:00 2001
From: skyleaworlder <870033938@qq.com>
Date: Fri, 23 Dec 2022 07:56:14 +0000
Subject: [PATCH 01/12] add: init my lppool impl

---
 src/NNlib.jl               |  4 ++--
 src/impl/pooling_direct.jl | 38 +++++++++++++++++++++++++++++++-------
 src/pooling.jl             | 14 +++++++++++---
 3 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/src/NNlib.jl b/src/NNlib.jl
index 19658d2be..f26d1f081 100644
--- a/src/NNlib.jl
+++ b/src/NNlib.jl
@@ -67,8 +67,8 @@ include("ctc.jl")
 export ctc_loss
 
 include("pooling.jl")
-export maxpool, maxpool!, meanpool, meanpool!, 
-    ∇maxpool, ∇maxpool!, ∇meanpool, ∇meanpool!
+export maxpool, maxpool!, meanpool, meanpool!, lppool, lppool!,
+    ∇maxpool, ∇maxpool!, ∇meanpool, ∇meanpool!, ∇lppool, ∇lppool!
 
 include("padding.jl")
 export pad_constant, pad_repeat, pad_reflect, pad_zeros
diff --git a/src/impl/pooling_direct.jl b/src/impl/pooling_direct.jl
index 566406eb2..55410e1de 100644
--- a/src/impl/pooling_direct.jl
+++ b/src/impl/pooling_direct.jl
@@ -1,14 +1,14 @@
 # Pooling is so similar, we abstract over meanpooling and maxpooling, simply replacing
 # the inner loop operation and a few initialization parameters.
-for name in (:max, :mean)
+for name in (:max, :mean, :lp)
     @eval function $((Symbol("$(name)pool_direct!")))(
                     y::AbstractArray{T, 5}, x::AbstractArray{T, 5},
-                    pdims::PoolDims; alpha::T=T(1), beta::T=T(0)) where T
+                    pdims::PoolDims; alpha::T=T(1), beta::T=T(0), kwargs...) where T
         $((Symbol("$(name)pool_direct!")))(
             y, x, pdims,
             Val(kernel_size(pdims)), Val(channels_out(pdims)),
             Val(padding(pdims)), Val(dilation(pdims)), Val(stride(pdims));
-            alpha, beta)
+            alpha, beta, kwargs...)
         return y
     end
 
@@ -17,7 +17,7 @@ for name in (:max, :mean)
         pdims::PoolDims,
         # kernel size, channels out, padding, dilation, stride
         ::Val{K}, ::Val{C}, ::Val{P}, ::Val{D}, ::Val{S};
-        alpha::T=T(1), beta::T=T(0),
+        alpha::T=T(1), beta::T=T(0), kwargs...
     ) where {T, K, C, P, D, S}
         @assert beta == T(0) "beta not supported yet"
         check_dims(size(x), size(y), pdims)
@@ -41,10 +41,15 @@ for name in (:max, :mean)
             alpha = alpha / prod(K)
         end
 
+        p = if $(name != :lp) 0 else
+            !haskey(kwargs, :p) && error("lppool must pass p")
+            kwargs[:p]
+        end
+
         # Each loop, we initialize `m` to something, set that here.
         m_init = if $(name == :max)
             T <: AbstractFloat ? nextfloat(typemin(T)) : typemin(T)
-        elseif $(name == :mean)
+        elseif $(name == :mean) || $(name == :lp)
             T(0)
         else
             error("Unimplemented codegen path")
@@ -78,11 +83,15 @@ for name in (:max, :mean)
                     end
                 elseif $(name == :mean)
                     m += x[input_kw, input_kh, input_kd, c, batch_idx]
+                elseif $(name == :lp)
+                    m += x[input_kw, input_kh, input_kd, c, batch_idx]^p
                 else
                     error("Unimplemented codegen path")
                 end
             end
 
+            m = $(name == :lp) ? m^(1 / p) : m
+
             y[w, h, d, c, batch_idx] = alpha * m # + beta * y[w, h, d, c, batch_idx]
             end
             end
@@ -128,12 +137,15 @@ for name in (:max, :mean)
                                 end
                             elseif $(name == :mean)
                                 m += x[input_kw, input_kh, input_kd, c, batch_idx]
+                            elseif $(name == :lp)
+                                m += x[input_kw, input_kh, input_kd, c, batch_idx]^p
                             else
                                 error("Unimplemented codegen path")
                             end
                         end
                     end
                 end
+                $(name == :lp) && (m = m^(1 / p))
                 y[w, h, d, c, batch_idx] = alpha * m # + beta * y[w, h, d, c, batch_idx]
                 end
                 end
@@ -159,7 +171,7 @@ for name in (:max, :mean)
                     dx::AbstractArray{T,5}, dy::AbstractArray{T,5},
                     y::AbstractArray{T,5}, x::AbstractArray{T,5},
                     pdims::PoolDims, ::Val{K}; # == kernel_size(pdims)
-                    alpha::T=T(1), beta::T=T(0)) where {T, K}
+                    alpha::T=T(1), beta::T=T(0), kwargs...) where {T, K}
         check_dims(size(x), size(dy), pdims)
 
         width, height, depth = input_size(pdims)
@@ -178,12 +190,20 @@ for name in (:max, :mean)
 
         # If we're doing mean pooling, we represent division by kernel size by rolling
         # it into the `alpha` multiplier.
-        if $(name == :mean)
+        if $(name == :mean) || $(name == :lp)
             alpha = alpha / prod(K)
         end
 
+        p = if $(name != :lp) 0 else
+            !haskey(kwargs, :p) && error("lppool must pass p")
+            kwargs[:p]
+        end
+
         # Start with the central region
         w_region, h_region, d_region = central_region
+        display(d_region)
+        display(h_region)
+        display(w_region)
         @inbounds for batch_idx in 1:size(x, 5), c in 1:out_c
             for d in d_region
             pd = project(d, stride_d, pad_d_lo)
@@ -226,6 +246,8 @@ for name in (:max, :mean)
                 elseif $(name == :mean)
                     # Either does meanpool :(
                     dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * alpha
+                elseif $(name == :lp)
+                    dx[input_kw, input_kh, input_kd, c, batch_idx] += (dy_idx^p * alpha)^(1/p)
                 else
                     error("Unimplemented codegen path")
                 end
@@ -286,6 +308,8 @@ for name in (:max, :mean)
                                 end
                             elseif $(name == :mean)
                                 dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * alpha #+ beta * dx[x_idxs...]
+                            elseif $(name == :lp)
+                                dx[input_kw, input_kh, input_kd, c, batch_idx] += (dy_idx^p * alpha)^(1/p)
                             else
                                 error("Unimplemented codegen path")
                             end
diff --git a/src/pooling.jl b/src/pooling.jl
index f13390342..c3c1f082a 100644
--- a/src/pooling.jl
+++ b/src/pooling.jl
@@ -26,6 +26,7 @@ for (front_name, backend) in (
         # This maps from public, front-facing name, to internal backend name
         :maxpool  => :direct,
         :meanpool => :direct,
+        :lppool => :direct,
     )
 
     # We only define 3d pooling primitives, we reshape lower down to get 1d and 2d pooling
@@ -42,6 +43,7 @@ end
 for (front_name, backend) in (
         :∇maxpool  => :direct,
         :∇meanpool => :direct,
+        :∇lppool => :direct,
     )
     @eval begin
         function $(Symbol("$(front_name)!"))(
@@ -57,7 +59,7 @@ end
 # Our strategy for pooling is to reshape to an array with three spatial dimensions, which
 # makes things MUCH EASIER for us on the backend side, and is in general pretty fast,
 # since we can specialize on sizes.
-for front_name in (:maxpool, :meanpool)
+for front_name in (:maxpool, :meanpool, :lppool)
     for backend in (Symbol(), :_direct)
         for N in (3, 4)
             @eval begin
@@ -103,7 +105,7 @@ end
 # Finally, let's generate auto-allocating versions of all our functions, for all backends:
 for backend in (Symbol(), :_direct, :_nnpack)
     # First make auto-allocating versions of the basic pooling calls:
-    for name in (:maxpool, :meanpool)
+    for name in (:maxpool, :meanpool, :lppool)
         @eval begin
             function $(Symbol("$(name)$(backend)"))(
                             x::AbstractArray{xT,N},
@@ -165,8 +167,14 @@ function meanpool(x, k::NTuple{N, Integer}; pad=0, stride=k) where N
     return meanpool(x, pdims)
 end
 
+function lppool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k) where N
+    pad = expand(Val(N), pad)
+    stride = expand(Val(N), stride)
+    pdims = PoolDims(x, k; padding=pad, stride=stride)
+    return lppool(x, pdims; p=p)
+end
 
-for pool in [:maxpool, :meanpool]
+for pool in [:maxpool, :meanpool, :lppool]
     ∇pool = Symbol(:∇, pool)
     pullback = Symbol(pool, :_pullback)
     @eval function rrule(::typeof($pool), x, pdims::PoolDims; kw...)

From ccd53cccb863e7bf8239776733af68a0c50a3c8c Mon Sep 17 00:00:00 2001
From: skyleaworlder <870033938@qq.com>
Date: Fri, 23 Dec 2022 15:00:18 +0000
Subject: [PATCH 02/12] add: impl gradient calculate

---
 src/impl/pooling_direct.jl | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/src/impl/pooling_direct.jl b/src/impl/pooling_direct.jl
index 55410e1de..c9088d024 100644
--- a/src/impl/pooling_direct.jl
+++ b/src/impl/pooling_direct.jl
@@ -201,9 +201,6 @@ for name in (:max, :mean, :lp)
 
         # Start with the central region
         w_region, h_region, d_region = central_region
-        display(d_region)
-        display(h_region)
-        display(w_region)
         @inbounds for batch_idx in 1:size(x, 5), c in 1:out_c
             for d in d_region
             pd = project(d, stride_d, pad_d_lo)
@@ -247,7 +244,8 @@ for name in (:max, :mean, :lp)
                     # Either does meanpool :(
                     dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * alpha
                 elseif $(name == :lp)
-                    dx[input_kw, input_kh, input_kd, c, batch_idx] += (dy_idx^p * alpha)^(1/p)
+                    grad = x[input_kw, input_kh, input_kd, c, batch_idx]^(p-1) * y_idx^(1-p)
+                    dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * grad
                 else
                     error("Unimplemented codegen path")
                 end
@@ -309,7 +307,8 @@ for name in (:max, :mean, :lp)
                             elseif $(name == :mean)
                                 dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * alpha #+ beta * dx[x_idxs...]
                             elseif $(name == :lp)
-                                dx[input_kw, input_kh, input_kd, c, batch_idx] += (dy_idx^p * alpha)^(1/p)
+                                grad = x[input_kw, input_kh, input_kd, c, batch_idx]^(p-1) * y_idx^(1-p)
+                                dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * grad
                             else
                                 error("Unimplemented codegen path")
                             end

From 07c2d022c81006b390e3fd16f859fb4049dec65f Mon Sep 17 00:00:00 2001
From: skyleaworlder <870033938@qq.com>
Date: Mon, 26 Dec 2022 14:35:21 +0000
Subject: [PATCH 03/12] test: for lppool(1d & 2d)

---
 test/pooling.jl | 106 ++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 106 insertions(+)

diff --git a/test/pooling.jl b/test/pooling.jl
index d1d26c620..ee525681f 100644
--- a/test/pooling.jl
+++ b/test/pooling.jl
@@ -248,6 +248,70 @@ meanpool_answer_dict = Dict(
     )
 )
 
+lppool_answer_dict = Dict(
+    1 => Dict(
+        "y"           => [2.019312856150994, 4.221163518110637],
+        "y_nostride"  => [
+            2.080083823051904, 3.2710663101885897,
+            4.497941445275415, 5.738793548317167
+        ],
+        "y_pad"       => [1.0, 3.605551275463989, 6.4031242374328485],
+        "dx"          => [
+            0.17258020254042603, 1.9525221042381296,
+            1.2774501198988355, 3.496467771732918, 0.0
+        ],
+        "dx_nostride" => [
+            0.48074985676913606, 3.1458422620080637,
+            4.752311710531486, 6.345225258061685, 4.356316321455918
+        ],
+        "dx_pad"       => [1.0, 2.0, 3.0, 4.0, 5.0],
+        "p"           => 4.5,
+        "p_nostride"  => 3.0,
+        "p_pad"       => 2.0
+    ),
+    2 => Dict(
+        "y"           => [
+            8.71909  24.9703;
+            11.7336  28.3804
+        ],
+        "y_nostride"  => [
+            11.1128  23.134   35.5704;
+            13.4219  25.6082  38.0707;
+            15.8033  28.0907  40.5735;
+            18.2249  30.5795  43.0782
+        ],
+        "y_pad"       => [
+            1.0      11.3616  16.0;
+            3.19158  15.9662  21.3545;
+            5.56869  18.7771  23.7903
+        ],
+        "dx"          => [
+            0.33866   4.97727  7.30092  12.8076;
+            0.957876  6.27208  8.31879  14.0269;
+            1.51693   6.6057   8.79844  14.3351;
+            2.33547   7.8822   9.83293  15.5461;
+            0.0       0.0      0.0      0.0 
+        ],
+        "dx_nostride" => [
+            3.33359  19.9471  35.7329  23.8564;
+            9.89551  44.627   76.2257  50.0307;
+           13.231    50.9101  82.5686  53.2022;
+           16.4888   57.223   88.9133  56.3742;
+            9.54591  30.9869  46.8371  29.3524
+        ],
+        "dx_pad"      => [
+            1.0       2.30261  10.4791   16.0;
+            0.992125  2.0321    7.81903  12.075;
+            2.73398   2.83743   9.5512   13.9299;
+            2.43512   2.98652   9.0132   13.5608;
+            4.25398   3.8865   10.7099   15.4161
+        ],
+        "p"           => 2.5,
+        "p_nostride"  => 1.5,
+        "p_pad"       => 3.5
+    )
+)
+
 for rank in (1, 2, 3)
     @testset "pool$(rank)d" begin
         for (pool, ∇pool, answer_dict) in (
@@ -297,6 +361,48 @@ for rank in (1, 2, 3)
     end
 end
 
+for rank in (1, 2)
+    for (pool, ∇pool, answer_dict) in (
+            (lppool, ∇lppool, lppool_answer_dict),
+            (NNlib.lppool_direct, NNlib.∇lppool_direct, lppool_answer_dict),)
+        @testset "$(pool)$(rank)d" begin
+            y = answer_dict[rank]["y"]
+            y_nostride = answer_dict[rank]["y_nostride"]
+            y_pad = answer_dict[rank]["y_pad"]
+            dx = answer_dict[rank]["dx"]
+            dx_nostride = answer_dict[rank]["dx_nostride"]
+            dx_pad = answer_dict[rank]["dx_pad"]
+            p = answer_dict[rank]["p"]
+            p_nostride = answer_dict[rank]["p_nostride"]
+            p_pad = answer_dict[rank]["p_pad"]
+
+            x = reshape(Float64[1:prod(size(dx));], size(dx)..., 1, 1)
+
+            ddims(x) = dropdims(x, dims=(rank + 1, rank + 2))
+
+            @test pool(x, PoolDims(x, 1); p=p) ≈ x atol = 1e-3
+
+            # Test vanilla pooling
+            pdims = PoolDims(x, 2)
+            y_hat = pool(x, pdims; p=p)
+            @test ddims(y_hat) ≈ y atol = 1e-3
+            @test ddims(∇pool(y_hat, y_hat, x, pdims; p=p)) ≈ dx atol = 1e-3
+
+            # Strided pooling
+            pdims = PoolDims(x, 2; stride=1)
+            y_hat = pool(x, pdims; p=p_nostride)
+            @test ddims(y_hat) ≈ y_nostride atol = 1e-3
+            @test ddims(∇pool(y_hat, y_hat, x, pdims; p=p_nostride)) ≈ dx_nostride atol = 1e-3
+
+            # Padded pooling
+            pdims = PoolDims(x, 2; padding=1)
+            y_hat = pool(x, pdims; p=p_pad)
+            @test ddims(y_hat) ≈ y_pad atol = 1e-3
+            @test ddims(∇pool(y_hat, y_hat, x, pdims; p=p_pad)) ≈ dx_pad atol = 1e-3
+        end
+    end
+end
+
 @testset "Pooling - Check Sizes" begin
     x = rand(10, 10, 3, 10)
     @test size(maxpool(x, (2, 2))) == (5, 5, 3, 10)

From 867f37ca293fd60facf2485fb4a8fb6f35b85f8d Mon Sep 17 00:00:00 2001
From: skyleaworlder <870033938@qq.com>
Date: Mon, 26 Dec 2022 15:06:01 +0000
Subject: [PATCH 04/12] doc: juliadoc & formula comment

---
 src/impl/pooling_direct.jl |  3 +++
 src/pooling.jl             | 11 +++++++++++
 test/perf/perf_report.jl   |  1 +
 3 files changed, 15 insertions(+)

diff --git a/src/impl/pooling_direct.jl b/src/impl/pooling_direct.jl
index c9088d024..a2ce666cf 100644
--- a/src/impl/pooling_direct.jl
+++ b/src/impl/pooling_direct.jl
@@ -84,12 +84,14 @@ for name in (:max, :mean, :lp)
                 elseif $(name == :mean)
                     m += x[input_kw, input_kh, input_kd, c, batch_idx]
                 elseif $(name == :lp)
+                    # y = (∑ x^p)^(1/p), here to calculate (∑ x^p)
                     m += x[input_kw, input_kh, input_kd, c, batch_idx]^p
                 else
                     error("Unimplemented codegen path")
                 end
             end
 
+            # for lppool, y = (∑ x^p)^(1/p)
             m = $(name == :lp) ? m^(1 / p) : m
 
             y[w, h, d, c, batch_idx] = alpha * m # + beta * y[w, h, d, c, batch_idx]
@@ -244,6 +246,7 @@ for name in (:max, :mean, :lp)
                     # Either does meanpool :(
                     dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * alpha
                 elseif $(name == :lp)
+                    # y = (∑ x^p)^(1/p), ∂y/∂x = x^(p-1) × y^(1-p)
                     grad = x[input_kw, input_kh, input_kd, c, batch_idx]^(p-1) * y_idx^(1-p)
                     dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * grad
                 else
diff --git a/src/pooling.jl b/src/pooling.jl
index c3c1f082a..9babcdf30 100644
--- a/src/pooling.jl
+++ b/src/pooling.jl
@@ -8,11 +8,15 @@
 #     - maxpool!(y, x, pdims)
 #     - meanpool(x, pdims)
 #     - meanpool!(y, x, pdims)
+#     - lppool(x, pdims)
+#     - lppool!(y, x, pdims)
 #   - Pooling input backprop
 #     - ∇maxpool(dy, y, x, pdims)
 #     - ∇maxpool!(dx, dy, y, x, pdims)
 #     - ∇meanpool(dy, y, x, pdims)
 #     - ∇meanpool!(dx, dy, y, x pdims)
+#     - ∇lppool(dy, y, x, pdims)
+#     - ∇lppool!(dx, dy, y, x pdims)
 #
 #   All methods require a `PoolDims` object to define the dimensions and optional
 #   elements of the convolution (stride, dilation, etc...), which is easily constructable
@@ -167,6 +171,12 @@ function meanpool(x, k::NTuple{N, Integer}; pad=0, stride=k) where N
     return meanpool(x, pdims)
 end
 
+
+"""
+    lppool(x, p::Number, k::NTuple; pad=0, stride=k)
+
+Perform Lp pool operation with `p`-norm and `window size `k` on input tensor `x`
+"""
 function lppool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k) where N
     pad = expand(Val(N), pad)
     stride = expand(Val(N), stride)
@@ -174,6 +184,7 @@ function lppool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k) where N
     return lppool(x, pdims; p=p)
 end
 
+
 for pool in [:maxpool, :meanpool, :lppool]
     ∇pool = Symbol(:∇, pool)
     pullback = Symbol(pool, :_pullback)
diff --git a/test/perf/perf_report.jl b/test/perf/perf_report.jl
index b1010f416..50b213cd4 100644
--- a/test/perf/perf_report.jl
+++ b/test/perf/perf_report.jl
@@ -93,6 +93,7 @@ for rank in (2,),
     for (pool, ∇pool, name) in (
             (NNlib.maxpool!, NNlib.∇maxpool!, "maxpool"),
             (NNlib.meanpool!, NNlib.∇meanpool!, "meanpool"),
+            (NNlib.lppool!, NNlib.∇lppool!, "lppool"),
         )
 
         t_fwd  = @benchmark $(pool)( $y, $x, $pdims)

From bf94c176344e907a637b1b830e4620eca7aa2d97 Mon Sep 17 00:00:00 2001
From: skyleaworlder <870033938@qq.com>
Date: Wed, 28 Dec 2022 09:48:01 +0000
Subject: [PATCH 05/12] doc: update for meanpool/maxpool, add params details
 and warning for lppool-maxpool

---
 docs/src/reference.md |  3 ++-
 src/pooling.jl        | 53 +++++++++++++++++++++++++++++++++++++++----
 2 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/docs/src/reference.md b/docs/src/reference.md
index f4ae5fdbf..743bf6776 100644
--- a/docs/src/reference.md
+++ b/docs/src/reference.md
@@ -44,12 +44,13 @@ logsoftmax
 
 ## Pooling
 
-`Flux`'s `AdaptiveMaxPool`, `AdaptiveMeanPool`, `GlobalMaxPool`, `GlobalMeanPool`, `MaxPool`, and `MeanPool` use `NNlib.PoolDims`, `NNlib.maxpool`, and `NNlib.meanpool` as their backend.
+`Flux`'s `AdaptiveMaxPool`, `AdaptiveMeanPool`, `GlobalMaxPool`, `GlobalMeanPool`, `MaxPool`, `MeanPool` and `LPPool` use `NNlib.PoolDims`, `NNlib.maxpool`, `NNlib.meanpool` and `NNlib.lppool` as their backend.
 
 ```@docs
 PoolDims
 maxpool
 meanpool
+lppool
 ```
 
 ## Padding
diff --git a/src/pooling.jl b/src/pooling.jl
index 9babcdf30..f85b3b072 100644
--- a/src/pooling.jl
+++ b/src/pooling.jl
@@ -147,11 +147,19 @@ expand(N, i::Integer) = ntuple(_ -> i, N)
 
 
 """
-    maxpool(x, k::NTuple; pad=0, stride=k)
+    maxpool(x, k::NTuple{N, Integer}; pad=0, stride=k)
 
 Perform max pool operation with window size `k` on input tensor `x`.
+
+* `x` and `k`: Usually, ndim(x) ∈ [3, 5], length(k) ∈ [1, 3], s.t. ndim(x) == length(k) + 2
+* `pad`: See [`pad_zeros`](@ref) for details.
+* `stride`: Stride for each spatial axis. `k` as default if not present.
 """
 function maxpool(x, k::NTuple{N, Integer}; pad=0, stride=k) where N
+    ndims(x) == length(k) + 2 || error("maxpool expects ndims(x) == length(k)+2,
+                                        dimension of x is $(ndims(x)),
+                                        length of k need $(ndims(x) - 2),
+                                        but now it's $(length(k))")
     pad = expand(Val(N), pad)
     stride = expand(Val(N), stride)
     pdims = PoolDims(x, k; padding=pad, stride=stride)
@@ -160,11 +168,19 @@ end
 
 
 """
-    meanpool(x, k::NTuple; pad=0, stride=k)
+    meanpool(x, k::NTuple{N, Integer}; pad=0, stride=k)
 
 Perform mean pool operation with window size `k` on input tensor `x`.
+
+* `x` and `k`: Usually, ndim(x) ∈ [3, 5], length(k) ∈ [1, 3], s.t. ndim(x) == length(k) + 2
+* `pad`: See [`pad_zeros`](@ref) for details.
+* `stride`: Stride for each spatial axis. `k` as default if not present.
 """
 function meanpool(x, k::NTuple{N, Integer}; pad=0, stride=k) where N
+    ndims(x) == length(k) + 2 || error("meanpool expects ndims(x) == length(k)+2,
+                                        dimension of x is $(ndims(x)),
+                                        length of k need $(ndims(x) - 2),
+                                        but now it's $(length(k))")
     pad = expand(Val(N), pad)
     stride = expand(Val(N), stride)
     pdims = PoolDims(x, k; padding=pad, stride=stride)
@@ -173,11 +189,40 @@ end
 
 
 """
-    lppool(x, p::Number, k::NTuple; pad=0, stride=k)
+    lppool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k)
+
+Perform Lp pool operation with value of the Lp norm `p` and `window size `k` on input tensor `x`.
+
+* `x` and `k`: Usually, ndim(x) ∈ [3, 5], length(k) ∈ [1, 3], s.t. ndim(x) == length(k) + 2
+* `pad`: See [`pad_zeros`](@ref) for details.
+* `stride`: Stride for each spatial axis. `k` as default if not present.
+
+For each element `x` in (k × k) window, lppool computes `(∑ x^p)^(1 / p)` as output.
+
+* When p = 1, lppool(x, p, k) ./ prod(k) ≈ meanpool(x, k)
+* When p = 2, lppool(x, p, k).^2 ./ prod(k) ≈ meanpool(x.^2, k)
+
+!!! warning
+
+    Theoretically, when `p -> ∞`, lppool(x, p, k) ≈ maxpool(x, k).
+    But it's not correct in julia. Given a normal valuable `x`,
+    ```jldoctest
+    julia> x = 10
+    10
+
+    julia> ans^Inf
+    Inf
 
-Perform Lp pool operation with `p`-norm and `window size `k` on input tensor `x`
+    julia> ans^(1/Inf)
+    1.0
+    ```
+    Please use `meanpool` and `maxpool` directly when needed.
 """
 function lppool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k) where N
+    ndims(x) == length(k) + 2 || error("lppool expects ndims(x) == length(k)+2,
+                                        dimension of x is $(ndims(x)),
+                                        length of k need $(ndims(x) - 2),
+                                        but now it's $(length(k))")
     pad = expand(Val(N), pad)
     stride = expand(Val(N), stride)
     pdims = PoolDims(x, k; padding=pad, stride=stride)

From 671371d2768b997936c85756f27fa3b4b0fae48a Mon Sep 17 00:00:00 2001
From: skyleaworlder <870033938@qq.com>
Date: Wed, 28 Dec 2022 10:56:57 +0000
Subject: [PATCH 06/12] update: 1 => T(1); remove unnecessary if-condition;
 clear error msg and doc

---
 src/impl/pooling_direct.jl | 8 ++++----
 src/pooling.jl             | 4 ++--
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/impl/pooling_direct.jl b/src/impl/pooling_direct.jl
index a2ce666cf..fb752702e 100644
--- a/src/impl/pooling_direct.jl
+++ b/src/impl/pooling_direct.jl
@@ -42,7 +42,7 @@ for name in (:max, :mean, :lp)
         end
 
         p = if $(name != :lp) 0 else
-            !haskey(kwargs, :p) && error("lppool must pass p")
+            !haskey(kwargs, :p) && error("lppool needs keyword argument `p`")
             kwargs[:p]
         end
 
@@ -92,7 +92,7 @@ for name in (:max, :mean, :lp)
             end
 
             # for lppool, y = (∑ x^p)^(1/p)
-            m = $(name == :lp) ? m^(1 / p) : m
+            m = $(name == :lp) ? m^(T(1) / p) : m
 
             y[w, h, d, c, batch_idx] = alpha * m # + beta * y[w, h, d, c, batch_idx]
             end
@@ -147,7 +147,7 @@ for name in (:max, :mean, :lp)
                         end
                     end
                 end
-                $(name == :lp) && (m = m^(1 / p))
+                $(name == :lp) && (m = m^(T(1) / p))
                 y[w, h, d, c, batch_idx] = alpha * m # + beta * y[w, h, d, c, batch_idx]
                 end
                 end
@@ -192,7 +192,7 @@ for name in (:max, :mean, :lp)
 
         # If we're doing mean pooling, we represent division by kernel size by rolling
         # it into the `alpha` multiplier.
-        if $(name == :mean) || $(name == :lp)
+        if $(name == :mean)
             alpha = alpha / prod(K)
         end
 
diff --git a/src/pooling.jl b/src/pooling.jl
index f85b3b072..7163d7522 100644
--- a/src/pooling.jl
+++ b/src/pooling.jl
@@ -205,9 +205,9 @@ For each element `x` in (k × k) window, lppool computes `(∑ x^p)^(1 / p)` as
 !!! warning
 
     Theoretically, when `p -> ∞`, lppool(x, p, k) ≈ maxpool(x, k).
-    But it's not correct in julia. Given a normal valuable `x`,
+    But it's not correct in julia. Given an arbitrary Number `n`,
     ```jldoctest
-    julia> x = 10
+    julia> n = 10
     10
 
     julia> ans^Inf

From 909378cd8507fab54e39ff9d1ed7a0e1eea643d9 Mon Sep 17 00:00:00 2001
From: skyleaworlder <870033938@qq.com>
Date: Fri, 30 Dec 2022 07:42:38 +0000
Subject: [PATCH 07/12] update: move dim checker from exported pool function to
 PoolDims

---
 src/dim_helpers/PoolDims.jl |  6 ++++++
 src/pooling.jl              | 12 ------------
 2 files changed, 6 insertions(+), 12 deletions(-)

diff --git a/src/dim_helpers/PoolDims.jl b/src/dim_helpers/PoolDims.jl
index e7219ff72..75d56b8cd 100644
--- a/src/dim_helpers/PoolDims.jl
+++ b/src/dim_helpers/PoolDims.jl
@@ -25,6 +25,12 @@ function PoolDims(
     _check_kernel(k::NTuple, ::Int) = k
 
     kernel = _check_kernel(k, M - 2)
+    length(x_size) == length(kernel) + 2 || error(
+        "PoolDims expects ndim(x) == length(k)+2 or length(size(x)) == length(kernel)+2,
+        dimension of x_size is $(length(x_size)),
+        length of k need $(length(x_size) - 2),
+        but now it's $(length(kernel))"
+    )
     spdf_kernel = NTuple{M, Int}([kernel..., 1, 1])
 
     sstride, ppadding, ddilation = check_spdf(
diff --git a/src/pooling.jl b/src/pooling.jl
index 7163d7522..4185936cf 100644
--- a/src/pooling.jl
+++ b/src/pooling.jl
@@ -156,10 +156,6 @@ Perform max pool operation with window size `k` on input tensor `x`.
 * `stride`: Stride for each spatial axis. `k` as default if not present.
 """
 function maxpool(x, k::NTuple{N, Integer}; pad=0, stride=k) where N
-    ndims(x) == length(k) + 2 || error("maxpool expects ndims(x) == length(k)+2,
-                                        dimension of x is $(ndims(x)),
-                                        length of k need $(ndims(x) - 2),
-                                        but now it's $(length(k))")
     pad = expand(Val(N), pad)
     stride = expand(Val(N), stride)
     pdims = PoolDims(x, k; padding=pad, stride=stride)
@@ -177,10 +173,6 @@ Perform mean pool operation with window size `k` on input tensor `x`.
 * `stride`: Stride for each spatial axis. `k` as default if not present.
 """
 function meanpool(x, k::NTuple{N, Integer}; pad=0, stride=k) where N
-    ndims(x) == length(k) + 2 || error("meanpool expects ndims(x) == length(k)+2,
-                                        dimension of x is $(ndims(x)),
-                                        length of k need $(ndims(x) - 2),
-                                        but now it's $(length(k))")
     pad = expand(Val(N), pad)
     stride = expand(Val(N), stride)
     pdims = PoolDims(x, k; padding=pad, stride=stride)
@@ -219,10 +211,6 @@ For each element `x` in (k × k) window, lppool computes `(∑ x^p)^(1 / p)` as
     Please use `meanpool` and `maxpool` directly when needed.
 """
 function lppool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k) where N
-    ndims(x) == length(k) + 2 || error("lppool expects ndims(x) == length(k)+2,
-                                        dimension of x is $(ndims(x)),
-                                        length of k need $(ndims(x) - 2),
-                                        but now it's $(length(k))")
     pad = expand(Val(N), pad)
     stride = expand(Val(N), stride)
     pdims = PoolDims(x, k; padding=pad, stride=stride)

From e78bd1bb86cb1ab25324c325c23117b7453d0a0d Mon Sep 17 00:00:00 2001
From: skyleaworlder <870033938@qq.com>
Date: Fri, 30 Dec 2022 08:16:45 +0000
Subject: [PATCH 08/12] rename: lppool => normpool

---
 docs/src/reference.md      |  4 ++--
 src/NNlib.jl               |  4 ++--
 src/impl/pooling_direct.jl | 26 ++++++++++----------
 src/pooling.jl             | 49 +++++++++++++-------------------------
 test/perf/perf_report.jl   |  2 +-
 test/pooling.jl            |  6 ++---
 6 files changed, 38 insertions(+), 53 deletions(-)

diff --git a/docs/src/reference.md b/docs/src/reference.md
index 743bf6776..8162f3afb 100644
--- a/docs/src/reference.md
+++ b/docs/src/reference.md
@@ -44,13 +44,13 @@ logsoftmax
 
 ## Pooling
 
-`Flux`'s `AdaptiveMaxPool`, `AdaptiveMeanPool`, `GlobalMaxPool`, `GlobalMeanPool`, `MaxPool`, `MeanPool` and `LPPool` use `NNlib.PoolDims`, `NNlib.maxpool`, `NNlib.meanpool` and `NNlib.lppool` as their backend.
+`Flux`'s `AdaptiveMaxPool`, `AdaptiveMeanPool`, `GlobalMaxPool`, `GlobalMeanPool`, `MaxPool`, `MeanPool` and `normpool` use `NNlib.PoolDims`, `NNlib.maxpool`, `NNlib.meanpool` and `NNlib.normpool` as their backend.
 
 ```@docs
 PoolDims
 maxpool
 meanpool
-lppool
+normpool
 ```
 
 ## Padding
diff --git a/src/NNlib.jl b/src/NNlib.jl
index f26d1f081..9f18174a0 100644
--- a/src/NNlib.jl
+++ b/src/NNlib.jl
@@ -67,8 +67,8 @@ include("ctc.jl")
 export ctc_loss
 
 include("pooling.jl")
-export maxpool, maxpool!, meanpool, meanpool!, lppool, lppool!,
-    ∇maxpool, ∇maxpool!, ∇meanpool, ∇meanpool!, ∇lppool, ∇lppool!
+export maxpool, maxpool!, meanpool, meanpool!, normpool, normpool!,
+    ∇maxpool, ∇maxpool!, ∇meanpool, ∇meanpool!, ∇normpool, ∇normpool!
 
 include("padding.jl")
 export pad_constant, pad_repeat, pad_reflect, pad_zeros
diff --git a/src/impl/pooling_direct.jl b/src/impl/pooling_direct.jl
index fb752702e..8fe5cbb79 100644
--- a/src/impl/pooling_direct.jl
+++ b/src/impl/pooling_direct.jl
@@ -1,6 +1,6 @@
 # Pooling is so similar, we abstract over meanpooling and maxpooling, simply replacing
 # the inner loop operation and a few initialization parameters.
-for name in (:max, :mean, :lp)
+for name in (:max, :mean, :norm)
     @eval function $((Symbol("$(name)pool_direct!")))(
                     y::AbstractArray{T, 5}, x::AbstractArray{T, 5},
                     pdims::PoolDims; alpha::T=T(1), beta::T=T(0), kwargs...) where T
@@ -41,15 +41,15 @@ for name in (:max, :mean, :lp)
             alpha = alpha / prod(K)
         end
 
-        p = if $(name != :lp) 0 else
-            !haskey(kwargs, :p) && error("lppool needs keyword argument `p`")
+        p = if $(name != :norm) 0 else
+            !haskey(kwargs, :p) && error("normpool needs keyword argument `p`")
             kwargs[:p]
         end
 
         # Each loop, we initialize `m` to something, set that here.
         m_init = if $(name == :max)
             T <: AbstractFloat ? nextfloat(typemin(T)) : typemin(T)
-        elseif $(name == :mean) || $(name == :lp)
+        elseif $(name == :mean) || $(name == :norm)
             T(0)
         else
             error("Unimplemented codegen path")
@@ -83,7 +83,7 @@ for name in (:max, :mean, :lp)
                     end
                 elseif $(name == :mean)
                     m += x[input_kw, input_kh, input_kd, c, batch_idx]
-                elseif $(name == :lp)
+                elseif $(name == :norm)
                     # y = (∑ x^p)^(1/p), here to calculate (∑ x^p)
                     m += x[input_kw, input_kh, input_kd, c, batch_idx]^p
                 else
@@ -91,8 +91,8 @@ for name in (:max, :mean, :lp)
                 end
             end
 
-            # for lppool, y = (∑ x^p)^(1/p)
-            m = $(name == :lp) ? m^(T(1) / p) : m
+            # for normpool, y = (∑ x^p)^(1/p)
+            m = $(name == :norm) ? m^(T(1) / p) : m
 
             y[w, h, d, c, batch_idx] = alpha * m # + beta * y[w, h, d, c, batch_idx]
             end
@@ -139,7 +139,7 @@ for name in (:max, :mean, :lp)
                                 end
                             elseif $(name == :mean)
                                 m += x[input_kw, input_kh, input_kd, c, batch_idx]
-                            elseif $(name == :lp)
+                            elseif $(name == :norm)
                                 m += x[input_kw, input_kh, input_kd, c, batch_idx]^p
                             else
                                 error("Unimplemented codegen path")
@@ -147,7 +147,7 @@ for name in (:max, :mean, :lp)
                         end
                     end
                 end
-                $(name == :lp) && (m = m^(T(1) / p))
+                $(name == :norm) && (m = m^(T(1) / p))
                 y[w, h, d, c, batch_idx] = alpha * m # + beta * y[w, h, d, c, batch_idx]
                 end
                 end
@@ -196,8 +196,8 @@ for name in (:max, :mean, :lp)
             alpha = alpha / prod(K)
         end
 
-        p = if $(name != :lp) 0 else
-            !haskey(kwargs, :p) && error("lppool must pass p")
+        p = if $(name != :norm) 0 else
+            !haskey(kwargs, :p) && error("normpool must pass p")
             kwargs[:p]
         end
 
@@ -245,7 +245,7 @@ for name in (:max, :mean, :lp)
                 elseif $(name == :mean)
                     # Either does meanpool :(
                     dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * alpha
-                elseif $(name == :lp)
+                elseif $(name == :norm)
                     # y = (∑ x^p)^(1/p), ∂y/∂x = x^(p-1) × y^(1-p)
                     grad = x[input_kw, input_kh, input_kd, c, batch_idx]^(p-1) * y_idx^(1-p)
                     dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * grad
@@ -309,7 +309,7 @@ for name in (:max, :mean, :lp)
                                 end
                             elseif $(name == :mean)
                                 dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * alpha #+ beta * dx[x_idxs...]
-                            elseif $(name == :lp)
+                            elseif $(name == :norm)
                                 grad = x[input_kw, input_kh, input_kd, c, batch_idx]^(p-1) * y_idx^(1-p)
                                 dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * grad
                             else
diff --git a/src/pooling.jl b/src/pooling.jl
index 4185936cf..dbeac5b53 100644
--- a/src/pooling.jl
+++ b/src/pooling.jl
@@ -8,15 +8,15 @@
 #     - maxpool!(y, x, pdims)
 #     - meanpool(x, pdims)
 #     - meanpool!(y, x, pdims)
-#     - lppool(x, pdims)
-#     - lppool!(y, x, pdims)
+#     - normpool(x, pdims)
+#     - normpool!(y, x, pdims)
 #   - Pooling input backprop
 #     - ∇maxpool(dy, y, x, pdims)
 #     - ∇maxpool!(dx, dy, y, x, pdims)
 #     - ∇meanpool(dy, y, x, pdims)
 #     - ∇meanpool!(dx, dy, y, x pdims)
-#     - ∇lppool(dy, y, x, pdims)
-#     - ∇lppool!(dx, dy, y, x pdims)
+#     - ∇normpool(dy, y, x, pdims)
+#     - ∇normpool!(dx, dy, y, x pdims)
 #
 #   All methods require a `PoolDims` object to define the dimensions and optional
 #   elements of the convolution (stride, dilation, etc...), which is easily constructable
@@ -30,7 +30,7 @@ for (front_name, backend) in (
         # This maps from public, front-facing name, to internal backend name
         :maxpool  => :direct,
         :meanpool => :direct,
-        :lppool => :direct,
+        :normpool => :direct,
     )
 
     # We only define 3d pooling primitives, we reshape lower down to get 1d and 2d pooling
@@ -47,7 +47,7 @@ end
 for (front_name, backend) in (
         :∇maxpool  => :direct,
         :∇meanpool => :direct,
-        :∇lppool => :direct,
+        :∇normpool => :direct,
     )
     @eval begin
         function $(Symbol("$(front_name)!"))(
@@ -63,7 +63,7 @@ end
 # Our strategy for pooling is to reshape to an array with three spatial dimensions, which
 # makes things MUCH EASIER for us on the backend side, and is in general pretty fast,
 # since we can specialize on sizes.
-for front_name in (:maxpool, :meanpool, :lppool)
+for front_name in (:maxpool, :meanpool, :normpool)
     for backend in (Symbol(), :_direct)
         for N in (3, 4)
             @eval begin
@@ -109,7 +109,7 @@ end
 # Finally, let's generate auto-allocating versions of all our functions, for all backends:
 for backend in (Symbol(), :_direct, :_nnpack)
     # First make auto-allocating versions of the basic pooling calls:
-    for name in (:maxpool, :meanpool, :lppool)
+    for name in (:maxpool, :meanpool, :normpool)
         @eval begin
             function $(Symbol("$(name)$(backend)"))(
                             x::AbstractArray{xT,N},
@@ -181,44 +181,29 @@ end
 
 
 """
-    lppool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k)
+    normpool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k)
 
-Perform Lp pool operation with value of the Lp norm `p` and `window size `k` on input tensor `x`.
+Perform Lp pool operation with value of the Lp norm `p` and window size `k` on input tensor `x`, also known as LPPool in pytorch.
 
 * `x` and `k`: Usually, ndim(x) ∈ [3, 5], length(k) ∈ [1, 3], s.t. ndim(x) == length(k) + 2
+* `p` is restricted to (0, Inf). Out-of-range `p` leads to undefined behavior.
 * `pad`: See [`pad_zeros`](@ref) for details.
 * `stride`: Stride for each spatial axis. `k` as default if not present.
 
-For each element `x` in (k × k) window, lppool computes `(∑ x^p)^(1 / p)` as output.
+For each element `x` in (k × k) window, normpool computes `(∑ x^p)^(1 / p)` as output.
 
-* When p = 1, lppool(x, p, k) ./ prod(k) ≈ meanpool(x, k)
-* When p = 2, lppool(x, p, k).^2 ./ prod(k) ≈ meanpool(x.^2, k)
-
-!!! warning
-
-    Theoretically, when `p -> ∞`, lppool(x, p, k) ≈ maxpool(x, k).
-    But it's not correct in julia. Given an arbitrary Number `n`,
-    ```jldoctest
-    julia> n = 10
-    10
-
-    julia> ans^Inf
-    Inf
-
-    julia> ans^(1/Inf)
-    1.0
-    ```
-    Please use `meanpool` and `maxpool` directly when needed.
+* When p = 1, normpool(x, p, k) ./ prod(k) ≈ meanpool(x, k)
+* When p = 2, normpool(x, p, k).^2 ./ prod(k) ≈ meanpool(x.^2, k)
 """
-function lppool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k) where N
+function normpool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k) where N
     pad = expand(Val(N), pad)
     stride = expand(Val(N), stride)
     pdims = PoolDims(x, k; padding=pad, stride=stride)
-    return lppool(x, pdims; p=p)
+    return normpool(x, pdims; p=p)
 end
 
 
-for pool in [:maxpool, :meanpool, :lppool]
+for pool in [:maxpool, :meanpool, :normpool]
     ∇pool = Symbol(:∇, pool)
     pullback = Symbol(pool, :_pullback)
     @eval function rrule(::typeof($pool), x, pdims::PoolDims; kw...)
diff --git a/test/perf/perf_report.jl b/test/perf/perf_report.jl
index 50b213cd4..badd40b3c 100644
--- a/test/perf/perf_report.jl
+++ b/test/perf/perf_report.jl
@@ -93,7 +93,7 @@ for rank in (2,),
     for (pool, ∇pool, name) in (
             (NNlib.maxpool!, NNlib.∇maxpool!, "maxpool"),
             (NNlib.meanpool!, NNlib.∇meanpool!, "meanpool"),
-            (NNlib.lppool!, NNlib.∇lppool!, "lppool"),
+            (NNlib.normpool!, NNlib.∇normpool!, "normpool"),
         )
 
         t_fwd  = @benchmark $(pool)( $y, $x, $pdims)
diff --git a/test/pooling.jl b/test/pooling.jl
index ee525681f..38b4d4ca0 100644
--- a/test/pooling.jl
+++ b/test/pooling.jl
@@ -248,7 +248,7 @@ meanpool_answer_dict = Dict(
     )
 )
 
-lppool_answer_dict = Dict(
+normpool_answer_dict = Dict(
     1 => Dict(
         "y"           => [2.019312856150994, 4.221163518110637],
         "y_nostride"  => [
@@ -363,8 +363,8 @@ end
 
 for rank in (1, 2)
     for (pool, ∇pool, answer_dict) in (
-            (lppool, ∇lppool, lppool_answer_dict),
-            (NNlib.lppool_direct, NNlib.∇lppool_direct, lppool_answer_dict),)
+            (normpool, ∇normpool, normpool_answer_dict),
+            (NNlib.normpool_direct, NNlib.∇normpool_direct, normpool_answer_dict),)
         @testset "$(pool)$(rank)d" begin
             y = answer_dict[rank]["y"]
             y_nostride = answer_dict[rank]["y_nostride"]

From 7572ea09c81cd90f30a801bb79233c08564bb905 Mon Sep 17 00:00:00 2001
From: skyleaworlder <870033938@qq.com>
Date: Fri, 30 Dec 2022 08:28:32 +0000
Subject: [PATCH 09/12] add: normpool p value checker, p must be in (0, Inf)

---
 src/pooling.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/pooling.jl b/src/pooling.jl
index dbeac5b53..ead1f088a 100644
--- a/src/pooling.jl
+++ b/src/pooling.jl
@@ -196,6 +196,7 @@ For each element `x` in (k × k) window, normpool computes `(∑ x^p)^(1 / p)` a
 * When p = 2, normpool(x, p, k).^2 ./ prod(k) ≈ meanpool(x.^2, k)
 """
 function normpool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k) where N
+    (isinf(p) || p < 0) && error("p value of Lp norm pool expect to be in (0, Inf), but p is $(p) now.")
     pad = expand(Val(N), pad)
     stride = expand(Val(N), stride)
     pdims = PoolDims(x, k; padding=pad, stride=stride)

From d075827aa127629a4767ce6cc33e87b52967a0b4 Mon Sep 17 00:00:00 2001
From: skyleaworlder <870033938@qq.com>
Date: Wed, 4 Jan 2023 08:18:07 +0000
Subject: [PATCH 10/12] rename: normpool => lpnormpool

---
 docs/src/reference.md      |  4 ++--
 src/NNlib.jl               |  4 ++--
 src/impl/pooling_direct.jl | 26 +++++++++++++-------------
 src/pooling.jl             | 30 +++++++++++++++---------------
 test/perf/perf_report.jl   |  2 +-
 test/pooling.jl            |  6 +++---
 6 files changed, 36 insertions(+), 36 deletions(-)

diff --git a/docs/src/reference.md b/docs/src/reference.md
index 8162f3afb..c790c806a 100644
--- a/docs/src/reference.md
+++ b/docs/src/reference.md
@@ -44,13 +44,13 @@ logsoftmax
 
 ## Pooling
 
-`Flux`'s `AdaptiveMaxPool`, `AdaptiveMeanPool`, `GlobalMaxPool`, `GlobalMeanPool`, `MaxPool`, `MeanPool` and `normpool` use `NNlib.PoolDims`, `NNlib.maxpool`, `NNlib.meanpool` and `NNlib.normpool` as their backend.
+`Flux`'s `AdaptiveMaxPool`, `AdaptiveMeanPool`, `GlobalMaxPool`, `GlobalMeanPool`, `MaxPool`, `MeanPool` and `lpnormpool` use `NNlib.PoolDims`, `NNlib.maxpool`, `NNlib.meanpool` and `NNlib.lpnormpool` as their backend.
 
 ```@docs
 PoolDims
 maxpool
 meanpool
-normpool
+lpnormpool
 ```
 
 ## Padding
diff --git a/src/NNlib.jl b/src/NNlib.jl
index 9f18174a0..c406cb8a6 100644
--- a/src/NNlib.jl
+++ b/src/NNlib.jl
@@ -67,8 +67,8 @@ include("ctc.jl")
 export ctc_loss
 
 include("pooling.jl")
-export maxpool, maxpool!, meanpool, meanpool!, normpool, normpool!,
-    ∇maxpool, ∇maxpool!, ∇meanpool, ∇meanpool!, ∇normpool, ∇normpool!
+export maxpool, maxpool!, meanpool, meanpool!, lpnormpool, lpnormpool!,
+    ∇maxpool, ∇maxpool!, ∇meanpool, ∇meanpool!, ∇lpnormpool, ∇lpnormpool!
 
 include("padding.jl")
 export pad_constant, pad_repeat, pad_reflect, pad_zeros
diff --git a/src/impl/pooling_direct.jl b/src/impl/pooling_direct.jl
index 8fe5cbb79..d421d7bbc 100644
--- a/src/impl/pooling_direct.jl
+++ b/src/impl/pooling_direct.jl
@@ -1,6 +1,6 @@
 # Pooling is so similar, we abstract over meanpooling and maxpooling, simply replacing
 # the inner loop operation and a few initialization parameters.
-for name in (:max, :mean, :norm)
+for name in (:max, :mean, :lpnorm)
     @eval function $((Symbol("$(name)pool_direct!")))(
                     y::AbstractArray{T, 5}, x::AbstractArray{T, 5},
                     pdims::PoolDims; alpha::T=T(1), beta::T=T(0), kwargs...) where T
@@ -41,15 +41,15 @@ for name in (:max, :mean, :norm)
             alpha = alpha / prod(K)
         end
 
-        p = if $(name != :norm) 0 else
-            !haskey(kwargs, :p) && error("normpool needs keyword argument `p`")
+        p = if $(name != :lpnorm) 0 else
+            !haskey(kwargs, :p) && error("lpnormpool needs keyword argument `p`")
             kwargs[:p]
         end
 
         # Each loop, we initialize `m` to something, set that here.
         m_init = if $(name == :max)
             T <: AbstractFloat ? nextfloat(typemin(T)) : typemin(T)
-        elseif $(name == :mean) || $(name == :norm)
+        elseif $(name == :mean) || $(name == :lpnorm)
             T(0)
         else
             error("Unimplemented codegen path")
@@ -83,7 +83,7 @@ for name in (:max, :mean, :norm)
                     end
                 elseif $(name == :mean)
                     m += x[input_kw, input_kh, input_kd, c, batch_idx]
-                elseif $(name == :norm)
+                elseif $(name == :lpnorm)
                     # y = (∑ x^p)^(1/p), here to calculate (∑ x^p)
                     m += x[input_kw, input_kh, input_kd, c, batch_idx]^p
                 else
@@ -91,8 +91,8 @@ for name in (:max, :mean, :norm)
                 end
             end
 
-            # for normpool, y = (∑ x^p)^(1/p)
-            m = $(name == :norm) ? m^(T(1) / p) : m
+            # for lpnormpool, y = (∑ x^p)^(1/p)
+            m = $(name == :lpnorm) ? m^(T(1) / p) : m
 
             y[w, h, d, c, batch_idx] = alpha * m # + beta * y[w, h, d, c, batch_idx]
             end
@@ -139,7 +139,7 @@ for name in (:max, :mean, :norm)
                                 end
                             elseif $(name == :mean)
                                 m += x[input_kw, input_kh, input_kd, c, batch_idx]
-                            elseif $(name == :norm)
+                            elseif $(name == :lpnorm)
                                 m += x[input_kw, input_kh, input_kd, c, batch_idx]^p
                             else
                                 error("Unimplemented codegen path")
@@ -147,7 +147,7 @@ for name in (:max, :mean, :norm)
                         end
                     end
                 end
-                $(name == :norm) && (m = m^(T(1) / p))
+                $(name == :lpnorm) && (m = m^(T(1) / p))
                 y[w, h, d, c, batch_idx] = alpha * m # + beta * y[w, h, d, c, batch_idx]
                 end
                 end
@@ -196,8 +196,8 @@ for name in (:max, :mean, :norm)
             alpha = alpha / prod(K)
         end
 
-        p = if $(name != :norm) 0 else
-            !haskey(kwargs, :p) && error("normpool must pass p")
+        p = if $(name != :lpnorm) 0 else
+            !haskey(kwargs, :p) && error("lpnormpool must pass p")
             kwargs[:p]
         end
 
@@ -245,7 +245,7 @@ for name in (:max, :mean, :norm)
                 elseif $(name == :mean)
                     # Either does meanpool :(
                     dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * alpha
-                elseif $(name == :norm)
+                elseif $(name == :lpnorm)
                     # y = (∑ x^p)^(1/p), ∂y/∂x = x^(p-1) × y^(1-p)
                     grad = x[input_kw, input_kh, input_kd, c, batch_idx]^(p-1) * y_idx^(1-p)
                     dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * grad
@@ -309,7 +309,7 @@ for name in (:max, :mean, :norm)
                                 end
                             elseif $(name == :mean)
                                 dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * alpha #+ beta * dx[x_idxs...]
-                            elseif $(name == :norm)
+                            elseif $(name == :lpnorm)
                                 grad = x[input_kw, input_kh, input_kd, c, batch_idx]^(p-1) * y_idx^(1-p)
                                 dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * grad
                             else
diff --git a/src/pooling.jl b/src/pooling.jl
index ead1f088a..3b08772ee 100644
--- a/src/pooling.jl
+++ b/src/pooling.jl
@@ -8,15 +8,15 @@
 #     - maxpool!(y, x, pdims)
 #     - meanpool(x, pdims)
 #     - meanpool!(y, x, pdims)
-#     - normpool(x, pdims)
-#     - normpool!(y, x, pdims)
+#     - lpnormpool(x, pdims)
+#     - lpnormpool!(y, x, pdims)
 #   - Pooling input backprop
 #     - ∇maxpool(dy, y, x, pdims)
 #     - ∇maxpool!(dx, dy, y, x, pdims)
 #     - ∇meanpool(dy, y, x, pdims)
 #     - ∇meanpool!(dx, dy, y, x pdims)
-#     - ∇normpool(dy, y, x, pdims)
-#     - ∇normpool!(dx, dy, y, x pdims)
+#     - ∇lpnormpool(dy, y, x, pdims)
+#     - ∇lpnormpool!(dx, dy, y, x pdims)
 #
 #   All methods require a `PoolDims` object to define the dimensions and optional
 #   elements of the convolution (stride, dilation, etc...), which is easily constructable
@@ -30,7 +30,7 @@ for (front_name, backend) in (
         # This maps from public, front-facing name, to internal backend name
         :maxpool  => :direct,
         :meanpool => :direct,
-        :normpool => :direct,
+        :lpnormpool => :direct,
     )
 
     # We only define 3d pooling primitives, we reshape lower down to get 1d and 2d pooling
@@ -47,7 +47,7 @@ end
 for (front_name, backend) in (
         :∇maxpool  => :direct,
         :∇meanpool => :direct,
-        :∇normpool => :direct,
+        :∇lpnormpool => :direct,
     )
     @eval begin
         function $(Symbol("$(front_name)!"))(
@@ -63,7 +63,7 @@ end
 # Our strategy for pooling is to reshape to an array with three spatial dimensions, which
 # makes things MUCH EASIER for us on the backend side, and is in general pretty fast,
 # since we can specialize on sizes.
-for front_name in (:maxpool, :meanpool, :normpool)
+for front_name in (:maxpool, :meanpool, :lpnormpool)
     for backend in (Symbol(), :_direct)
         for N in (3, 4)
             @eval begin
@@ -109,7 +109,7 @@ end
 # Finally, let's generate auto-allocating versions of all our functions, for all backends:
 for backend in (Symbol(), :_direct, :_nnpack)
     # First make auto-allocating versions of the basic pooling calls:
-    for name in (:maxpool, :meanpool, :normpool)
+    for name in (:maxpool, :meanpool, :lpnormpool)
         @eval begin
             function $(Symbol("$(name)$(backend)"))(
                             x::AbstractArray{xT,N},
@@ -181,7 +181,7 @@ end
 
 
 """
-    normpool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k)
+    lpnormpool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k)
 
 Perform Lp pool operation with value of the Lp norm `p` and window size `k` on input tensor `x`, also known as LPPool in pytorch.
 
@@ -190,21 +190,21 @@ Perform Lp pool operation with value of the Lp norm `p` and window size `k` on i
 * `pad`: See [`pad_zeros`](@ref) for details.
 * `stride`: Stride for each spatial axis. `k` as default if not present.
 
-For each element `x` in (k × k) window, normpool computes `(∑ x^p)^(1 / p)` as output.
+For each element `x` in (k × k) window, lpnormpool computes `(∑ x^p)^(1 / p)` as output.
 
-* When p = 1, normpool(x, p, k) ./ prod(k) ≈ meanpool(x, k)
-* When p = 2, normpool(x, p, k).^2 ./ prod(k) ≈ meanpool(x.^2, k)
+* When p = 1, lpnormpool(x, p, k) ./ prod(k) ≈ meanpool(x, k)
+* When p = 2, lpnormpool(x, p, k).^2 ./ prod(k) ≈ meanpool(x.^2, k)
 """
-function normpool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k) where N
+function lpnormpool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k) where N
     (isinf(p) || p < 0) && error("p value of Lp norm pool expect to be in (0, Inf), but p is $(p) now.")
     pad = expand(Val(N), pad)
     stride = expand(Val(N), stride)
     pdims = PoolDims(x, k; padding=pad, stride=stride)
-    return normpool(x, pdims; p=p)
+    return lpnormpool(x, pdims; p=p)
 end
 
 
-for pool in [:maxpool, :meanpool, :normpool]
+for pool in [:maxpool, :meanpool, :lpnormpool]
     ∇pool = Symbol(:∇, pool)
     pullback = Symbol(pool, :_pullback)
     @eval function rrule(::typeof($pool), x, pdims::PoolDims; kw...)
diff --git a/test/perf/perf_report.jl b/test/perf/perf_report.jl
index badd40b3c..5c06515eb 100644
--- a/test/perf/perf_report.jl
+++ b/test/perf/perf_report.jl
@@ -93,7 +93,7 @@ for rank in (2,),
     for (pool, ∇pool, name) in (
             (NNlib.maxpool!, NNlib.∇maxpool!, "maxpool"),
             (NNlib.meanpool!, NNlib.∇meanpool!, "meanpool"),
-            (NNlib.normpool!, NNlib.∇normpool!, "normpool"),
+            (NNlib.lpnormpool!, NNlib.∇lpnormpool!, "lpnormpool"),
         )
 
         t_fwd  = @benchmark $(pool)( $y, $x, $pdims)
diff --git a/test/pooling.jl b/test/pooling.jl
index 38b4d4ca0..b7952295f 100644
--- a/test/pooling.jl
+++ b/test/pooling.jl
@@ -248,7 +248,7 @@ meanpool_answer_dict = Dict(
     )
 )
 
-normpool_answer_dict = Dict(
+lpnormpool_answer_dict = Dict(
     1 => Dict(
         "y"           => [2.019312856150994, 4.221163518110637],
         "y_nostride"  => [
@@ -363,8 +363,8 @@ end
 
 for rank in (1, 2)
     for (pool, ∇pool, answer_dict) in (
-            (normpool, ∇normpool, normpool_answer_dict),
-            (NNlib.normpool_direct, NNlib.∇normpool_direct, normpool_answer_dict),)
+            (lpnormpool, ∇lpnormpool, lpnormpool_answer_dict),
+            (NNlib.lpnormpool_direct, NNlib.∇lpnormpool_direct, lpnormpool_answer_dict),)
         @testset "$(pool)$(rank)d" begin
             y = answer_dict[rank]["y"]
             y_nostride = answer_dict[rank]["y_nostride"]

From 92f0327adc228fbd930465a92eb918ea0dc73dfe Mon Sep 17 00:00:00 2001
From: skyleaworlder <870033938@qq.com>
Date: Wed, 4 Jan 2023 08:28:56 +0000
Subject: [PATCH 11/12] doc: add paper reference to lp pooling

---
 src/pooling.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/pooling.jl b/src/pooling.jl
index 3b08772ee..9b01a499f 100644
--- a/src/pooling.jl
+++ b/src/pooling.jl
@@ -184,6 +184,7 @@ end
     lpnormpool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k)
 
 Perform Lp pool operation with value of the Lp norm `p` and window size `k` on input tensor `x`, also known as LPPool in pytorch.
+This pooling operator from [Learned-Norm Pooling for Deep Feedforward and Recurrent Neural Networks](https://arxiv.org/abs/1311.1780).
 
 * `x` and `k`: Usually, ndim(x) ∈ [3, 5], length(k) ∈ [1, 3], s.t. ndim(x) == length(k) + 2
 * `p` is restricted to (0, Inf). Out-of-range `p` leads to undefined behavior.

From 1f6e68fff879dc846b0010e552db9cb92713df60 Mon Sep 17 00:00:00 2001
From: skyleaworlder <870033938@qq.com>
Date: Fri, 6 Jan 2023 14:44:29 +0000
Subject: [PATCH 12/12] doc: usage and parameter description of pool functions

---
 src/impl/pooling_direct.jl |  6 +++---
 src/pooling.jl             | 27 ++++++++++++++++-----------
 2 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/src/impl/pooling_direct.jl b/src/impl/pooling_direct.jl
index d421d7bbc..55c1b9b23 100644
--- a/src/impl/pooling_direct.jl
+++ b/src/impl/pooling_direct.jl
@@ -84,14 +84,14 @@ for name in (:max, :mean, :lpnorm)
                 elseif $(name == :mean)
                     m += x[input_kw, input_kh, input_kd, c, batch_idx]
                 elseif $(name == :lpnorm)
-                    # y = (∑ x^p)^(1/p), here to calculate (∑ x^p)
+                    # y = (∑ᵢ xᵢ^p)^(1 / p), here to calculate ∑ᵢ xᵢ^p
                     m += x[input_kw, input_kh, input_kd, c, batch_idx]^p
                 else
                     error("Unimplemented codegen path")
                 end
             end
 
-            # for lpnormpool, y = (∑ x^p)^(1/p)
+            # for lpnormpool, y = (∑ᵢ xᵢ^p)^(1 / p)
             m = $(name == :lpnorm) ? m^(T(1) / p) : m
 
             y[w, h, d, c, batch_idx] = alpha * m # + beta * y[w, h, d, c, batch_idx]
@@ -246,7 +246,7 @@ for name in (:max, :mean, :lpnorm)
                     # Either does meanpool :(
                     dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * alpha
                 elseif $(name == :lpnorm)
-                    # y = (∑ x^p)^(1/p), ∂y/∂x = x^(p-1) × y^(1-p)
+                    # y = (∑ᵢ xᵢ^p)^(1 / p), ∂y/∂xᵢ = xᵢ^(p-1) × y^(1-p)
                     grad = x[input_kw, input_kh, input_kd, c, batch_idx]^(p-1) * y_idx^(1-p)
                     dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * grad
                 else
diff --git a/src/pooling.jl b/src/pooling.jl
index 9b01a499f..704d7d0fa 100644
--- a/src/pooling.jl
+++ b/src/pooling.jl
@@ -151,9 +151,11 @@ expand(N, i::Integer) = ntuple(_ -> i, N)
 
 Perform max pool operation with window size `k` on input tensor `x`.
 
-* `x` and `k`: Usually, ndim(x) ∈ [3, 5], length(k) ∈ [1, 3], s.t. ndim(x) == length(k) + 2
+Arguments:
+
+* `x` and `k`: Expects `ndim(x) ∈ 3:5`, and always `length(k) == ndim(x) - 2`
 * `pad`: See [`pad_zeros`](@ref) for details.
-* `stride`: Stride for each spatial axis. `k` as default if not present.
+* `stride`: Either a tuple with the same length as `k`, or one integer for all directions. Default is `k`.
 """
 function maxpool(x, k::NTuple{N, Integer}; pad=0, stride=k) where N
     pad = expand(Val(N), pad)
@@ -168,9 +170,11 @@ end
 
 Perform mean pool operation with window size `k` on input tensor `x`.
 
-* `x` and `k`: Usually, ndim(x) ∈ [3, 5], length(k) ∈ [1, 3], s.t. ndim(x) == length(k) + 2
+Arguments:
+
+* `x` and `k`: Expects `ndim(x) ∈ 3:5``, and always `length(k) == ndim(x) - 2`
 * `pad`: See [`pad_zeros`](@ref) for details.
-* `stride`: Stride for each spatial axis. `k` as default if not present.
+* `stride`: Either a tuple with the same length as `k`, or one integer for all directions. Default is `k`.
 """
 function meanpool(x, k::NTuple{N, Integer}; pad=0, stride=k) where N
     pad = expand(Val(N), pad)
@@ -186,18 +190,19 @@ end
 Perform Lp pool operation with value of the Lp norm `p` and window size `k` on input tensor `x`, also known as LPPool in pytorch.
 This pooling operator from [Learned-Norm Pooling for Deep Feedforward and Recurrent Neural Networks](https://arxiv.org/abs/1311.1780).
 
-* `x` and `k`: Usually, ndim(x) ∈ [3, 5], length(k) ∈ [1, 3], s.t. ndim(x) == length(k) + 2
-* `p` is restricted to (0, Inf). Out-of-range `p` leads to undefined behavior.
+Arguments:
+
+* `x` and `k`: Expects `ndim(x) ∈ 3:5``, and always `length(k) == ndim(x) - 2`
+* `p` is restricted to `0 < p < Inf`.
 * `pad`: See [`pad_zeros`](@ref) for details.
-* `stride`: Stride for each spatial axis. `k` as default if not present.
+* `stride`: Either a tuple with the same length as `k`, or one integer for all directions. Default is `k`.
 
-For each element `x` in (k × k) window, lpnormpool computes `(∑ x^p)^(1 / p)` as output.
+For all elements `x` in a size `k` window, lpnormpool computes `(∑ᵢ xᵢ^p)^(1 / p)` as an element of the output.
 
-* When p = 1, lpnormpool(x, p, k) ./ prod(k) ≈ meanpool(x, k)
-* When p = 2, lpnormpool(x, p, k).^2 ./ prod(k) ≈ meanpool(x.^2, k)
+Thus `lpnormpool(x, 1, k) ./ prod(k) ≈ meanpool(x, k)` and `lpnormpool(x, 2, k).^2 ./ prod(k) ≈ meanpool(x.^2, k)`.
 """
 function lpnormpool(x, p::Number, k::NTuple{N, Integer}; pad=0, stride=k) where N
-    (isinf(p) || p < 0) && error("p value of Lp norm pool expect to be in (0, Inf), but p is $(p) now.")
+    (isinf(p) || p < 0) && error("p value of Lp norm pool expects `0 < p < Inf`, but p is $(p) now.")
     pad = expand(Val(N), pad)
     stride = expand(Val(N), stride)
     pdims = PoolDims(x, k; padding=pad, stride=stride)