Merge pull request #387 from vincentmolin/rruleconv

add `rrule(::typeof(∇conv_filter)`

Author: ToucheSir

src/conv.jl

@@ -329,6 +329,19 @@ for conv in [:conv, :depthwiseconv]
     end
 end
 
+function rrule(::typeof(∇conv_filter), x, dy, cdims; kw...)
+    function ∇conv_filter_pullback(Δ)
+        Δ1 = colmajor(unthunk(Δ))
+        return (
+            NoTangent(),
+            @thunk(∇conv_data(dy, Δ1, cdims, kw...)),
+            @thunk(conv(x, Δ1, cdims, kw...)),
+            NoTangent(),
+        )
+    end
+    return ∇conv_filter(x, dy, cdims; kw...), ∇conv_filter_pullback
+end
+
 # Use NNPACK if it is available and the operation is supported
 # commented out 'till proper benchmarking and more correctness test are performed
 # if is_nnpack_available()

test/conv.jl

@@ -734,6 +734,10 @@ end
   # else
     gradtest((y, w) -> sum(∇conv_data(y, w, cdims)), y, w)
   # end
+    gradtest((x, y) -> ∇conv_filter(x, y, cdims), x, y)
+    if spatial_rank < 3
+        gradtest((x, y) -> sum(∇conv_filter(x, y, cdims)), x, y)
+    end
 
   dcdims = DepthwiseConvDims(x, w)
   gradtest((x, w) -> depthwiseconv(x, w, dcdims), x, w)