From df035ec12e5c34094f48e1e40667070e37363aa0 Mon Sep 17 00:00:00 2001
From: Rajeshkumar K <57673410+IBArbitrary@users.noreply.github.com>
Date: Fri, 6 Jun 2025 17:37:32 +0530
Subject: [PATCH 1/2] Fix lpnormpool

Implemented the correct LpNorm Pooling and backprop steps
---
 src/impl/pooling_direct.jl | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/impl/pooling_direct.jl b/src/impl/pooling_direct.jl
index 78a24ac4d..130477c80 100644
--- a/src/impl/pooling_direct.jl
+++ b/src/impl/pooling_direct.jl
@@ -96,7 +96,7 @@ for name in (:max, :mean, :lpnorm)
                     m += x[input_kw, input_kh, input_kd, c, batch_idx]
                 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
+                    m += abs(x[input_kw, input_kh, input_kd, c, batch_idx])^p
                 else
                     error("Unimplemented codegen path")
                 end
@@ -151,7 +151,7 @@ for name in (:max, :mean, :lpnorm)
                             elseif $(name == :mean)
                                 m += x[input_kw, input_kh, input_kd, c, batch_idx]
                             elseif $(name == :lpnorm)
-                                m += x[input_kw, input_kh, input_kd, c, batch_idx]^p
+                                m += abs(x[input_kw, input_kh, input_kd, c, batch_idx])^p
                             else
                                 error("Unimplemented codegen path")
                             end
@@ -264,7 +264,8 @@ for name in (:max, :mean, :lpnorm)
                     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)
-                    grad = x[input_kw, input_kh, input_kd, c, batch_idx]^(p-1) * y_idx^(1-p)
+                    xv = x[input_kw, input_kh, input_kd, c, batch_idx]
+                    grad = abs(xv)^(p-1) * y_idx^(1-p) * sign(xv)
                     dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * grad
                 else
                     error("Unimplemented codegen path")
@@ -327,7 +328,8 @@ for name in (:max, :mean, :lpnorm)
                             elseif $(name == :mean)
                                 dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * _alpha #+ _beta * dx[x_idxs...]
                             elseif $(name == :lpnorm)
-                                grad = x[input_kw, input_kh, input_kd, c, batch_idx]^(p-1) * y_idx^(1-p)
+                                xv = x[input_kw, input_kh, input_kd, c, batch_idx]
+                                grad = abs(xv)^(p-1) * y_idx^(1-p) * sign(xv)
                                 dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * grad
                             else
                                 error("Unimplemented codegen path")

From fa86b16a3b7b70f17baa494387fc24adc597f459 Mon Sep 17 00:00:00 2001
From: Rajeshkumar K <57673410+IBArbitrary@users.noreply.github.com>
Date: Fri, 6 Jun 2025 18:02:57 +0530
Subject: [PATCH 2/2] Added comments to lpnormpool

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

diff --git a/src/impl/pooling_direct.jl b/src/impl/pooling_direct.jl
index 130477c80..a0e65e344 100644
--- a/src/impl/pooling_direct.jl
+++ b/src/impl/pooling_direct.jl
@@ -95,14 +95,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 += abs(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]
@@ -263,7 +263,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) × sign(xᵢ)
                     xv = x[input_kw, input_kh, input_kd, c, batch_idx]
                     grad = abs(xv)^(p-1) * y_idx^(1-p) * sign(xv)
                     dx[input_kw, input_kh, input_kd, c, batch_idx] += dy_idx * grad