Add unthunk in conv.jl (#620)

FluxML · Jan 9, 2025 · b56ff50 · b56ff50 · CarloLucibello · Jan 9, 2025
1 parent 5732b97
commit b56ff50
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "NNlib"
 uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
-version = "0.9.26"
+version = "0.9.27"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/src/conv.jl b/src/conv.jl
@@ -343,25 +343,25 @@ for conv in [:conv, :depthwiseconv]
     conv_pullback, ∇conv_data_pullback = Symbol.([conv, ∇conv_data], :_pullback)
 
     @eval function rrule(::typeof($conv), x, w, cdims; kw...)
-        function $conv_pullback(Δ)
-            Δ = colmajor(Δ)
+        function $conv_pullback(Δraw)
+            Δ = colmajor(unthunk(Δraw))
             return (
                 NoTangent(),
-                @thunk($∇conv_data(unthunk(Δ), w, cdims, kw...)),
-                @thunk($∇conv_filter(x, unthunk(Δ), cdims, kw...)),
+                @thunk($∇conv_data(Δ, w, cdims, kw...)),
+                @thunk($∇conv_filter(x, Δ, cdims, kw...)),
                 NoTangent(),
             )
         end
         return $conv(x, w, cdims; kw...), $conv_pullback
     end
 
     @eval function rrule(::typeof($∇conv_data), x, w, cdims; kw...)
-        function $∇conv_data_pullback(Δ)
-            Δ = colmajor(Δ)
+        function $∇conv_data_pullback(Δraw)
+            Δ = colmajor(unthunk(Δraw))
             return (
                 NoTangent(),
-                @thunk($conv(unthunk(Δ), w, cdims, kw...)),
-                @thunk($∇conv_filter(unthunk(Δ), x, cdims, kw...)),
+                @thunk($conv(Δ, w, cdims, kw...)),
+                @thunk($∇conv_filter(Δ, x, cdims, kw...)),
                 NoTangent(),
             )
         end