Merge pull request #281 from FluxML/cl/rrule

cleanup AD

Author: CarloLucibello

src/NNlib.jl

@@ -3,6 +3,7 @@ module NNlib
 using Pkg
 using Requires
 using ChainRulesCore
+import ChainRulesCore: rrule
 using Base.Broadcast: broadcasted
 
 const Numeric = Union{AbstractArray{<:T}, T} where {T<:Number}

src/activations.jl

@@ -256,8 +256,8 @@ for (f, df) in UNARY_ACTS
     @eval @scalar_rule($f(x), $df)
 
     pullback = Symbol(:broadcasted_, f, :_pullback)
-    @eval function ChainRulesCore.rrule(::typeof(broadcasted),
-                                        ::typeof($f), x::Numeric)
+    @eval function rrule(::typeof(broadcasted),
+                         ::typeof($f), x::Numeric)
         Ω = $f.(x)
         function $pullback(Δ) 
             NO_FIELDS, NO_FIELDS, @.(Δ * $df)
@@ -275,9 +275,9 @@ for (f, df1, df2) in BINARY_ACTS
     @eval @scalar_rule($f(x1, x2), ($df1, $df2))
 
     pullback = Symbol(:broadcasted_, f, :_pullback)
-    @eval function ChainRulesCore.rrule(::typeof(broadcasted),
-                                        ::typeof($f), 
-                                        x1::Numeric, x2::Numeric)
+    @eval function rrule(::typeof(broadcasted),
+                         ::typeof($f), 
+                         x1::Numeric, x2::Numeric)
         Ω = $f.(x1, x2)
         function $pullback(Δ) 
             NO_FIELDS, NO_FIELDS, @.(Δ * $df1), @.(Δ * $df2)

src/batched/batchedadjtrans.jl

@@ -93,11 +93,11 @@ Base.unsafe_convert(::Type{Ptr{T}}, A::BatchedAdjOrTrans{T}) where {T} =
     Base.unsafe_convert(Ptr{T}, parent(A))
 
 # Gradients
-function ChainRulesCore.rrule(::typeof(batched_transpose), A::AbstractArray{<:Any,3})
+function rrule(::typeof(batched_transpose), A::AbstractArray{<:Any,3})
     b_transpose_back(Δ) = (NO_FIELDS, batched_transpose(Δ))
     batched_transpose(A), b_transpose_back
 end
-function ChainRulesCore.rrule(::typeof(batched_adjoint), A::AbstractArray{<:Any,3})
+function rrule(::typeof(batched_adjoint), A::AbstractArray{<:Any,3})
     b_adjoint_back(Δ) = (NO_FIELDS, batched_adjoint(Δ))
     batched_adjoint(A), b_adjoint_back
 end

src/batched/batchedmul.jl

@@ -56,11 +56,11 @@ function batched_mul(A::AbstractArray{T1, 3}, B::AbstractArray{T2, 3}) where {T1
     _batched_mul(storage_typejoin(A, B), A, B)
 end
 
-function ChainRulesCore.rrule(
-    ::typeof(batched_mul),
-    A::AbstractArray{S,3},
-    B::AbstractArray{T,3},
-) where {S,T}
+function rrule(::typeof(batched_mul),
+               A::AbstractArray{S,3},
+               B::AbstractArray{T,3},
+              ) where {S,T}
+    
     function batched_mul_pullback(Δ)
         return (
             NO_FIELDS,

src/conv.jl

@@ -211,11 +211,7 @@ for front_name in (:conv, :∇conv_data, :∇conv_filter,
 end
 
 for Dims in [:DenseConvDims, :DepthwiseConvDims, :PoolDims]
-    pullback = Symbol(Dims, :_pullback)
-    @eval function ChainRulesCore.rrule(::Type{$Dims}, args...; kwargs...)
-        $pullback(Δ) = (NO_FIELDS, ntuple(_ -> DoesNotExist(), length(args))...)
-        return $Dims(args...; kwargs...), $pullback
-    end
+    @eval @non_differentiable $Dims(::Any...)
 end
 
 colmajor(x) = (is_strided(x) && Base.stride(x, 1) == 1) ? x : collect(x)
@@ -224,7 +220,7 @@ for conv in [:conv, :depthwiseconv]
     local ∇conv_data, ∇conv_filter = Symbol.(:∇, conv, [:_data, :_filter])
     conv_pullback, ∇conv_data_pullback = Symbol.([conv, ∇conv_data], :_pullback)
 
-    @eval function ChainRulesCore.rrule(::typeof($conv), x, w, cdims; kw...)
+    @eval function rrule(::typeof($conv), x, w, cdims; kw...)
         function $conv_pullback(Δ)
             Δ = colmajor(Δ)
             return (
@@ -237,7 +233,7 @@ for conv in [:conv, :depthwiseconv]
         return $conv(x, w, cdims; kw...), $conv_pullback
     end
 
-    @eval function ChainRulesCore.rrule(::typeof($∇conv_data), x, w, cdims; kw...)
+    @eval function rrule(::typeof($∇conv_data), x, w, cdims; kw...)
         function $∇conv_data_pullback(Δ)
             Δ = colmajor(Δ)
             return (

src/pooling.jl

@@ -171,7 +171,7 @@ end
 for pool in [:maxpool, :meanpool]
     ∇pool = Symbol(:∇, pool)
     pullback = Symbol(pool, :_pullback)
-    @eval function ChainRulesCore.rrule(::typeof($pool), x, pdims::PoolDims; kw...)
+    @eval function rrule(::typeof($pool), x, pdims::PoolDims; kw...)
         Ω = $pool(x, pdims; kw...)
         $pullback(Δ) = (NO_FIELDS, $∇pool(Δ, Ω, x, pdims; kw...), DoesNotExist())
         return Ω, $pullback

src/softmax.jl

@@ -76,7 +76,7 @@ end
 ∇softmax!(Δ, x; dims=1) = ∇softmax!(Δ, Δ, x, softmax(x, dims=dims); dims=dims)
 ∇softmax!(out, Δ, x; dims=1) = ∇softmax!(out, Δ, x, softmax(x, dims=dims); dims=dims)
 
-function ChainRulesCore.rrule(::typeof(softmax), xs; dims=1)
+function rrule(::typeof(softmax), xs; dims=1)
     y = softmax(xs; dims=dims)
     softmax_pullback(Δ) = (NO_FIELDS, ∇softmax(Δ, xs, y, dims=dims))
     return y, softmax_pullback
@@ -112,7 +112,7 @@ end
 
 ∇logsoftmax(Δ::AbstractArray{T}, x::AbstractArray, y::AbstractArray{S}; dims = 1) where {T,S} =
     ∇logsoftmax!(similar(y, promote_type(T, S)), Δ, x, y; dims = dims)
-    
+
 # Old 2-arg version recomputing forward
 ∇logsoftmax(Δ, x; dims=1) =  ∇logsoftmax(Δ, x, logsoftmax(x, dims=dims); dims=dims)
 ∇logsoftmax!(Δ, x; dims=1) =  ∇logsoftmax!(Δ, Δ, x, logsoftmax(x, dims=dims); dims=dims)
@@ -123,7 +123,7 @@ function ∇logsoftmax!(out::AbstractArray, Δ::AbstractArray,
     out .= Δ .- sum(Δ, dims = dims) .* exp.(y)
 end
 
-function ChainRulesCore.rrule(::typeof(logsoftmax), xs; dims=1)
+function rrule(::typeof(logsoftmax), xs; dims=1)
     y = logsoftmax(xs; dims=dims)
     logsoftmax_pullback(Δ) = (NO_FIELDS, ∇logsoftmax(Δ, xs, y, dims=dims))
     return y, logsoftmax_pullback

src/upsample.jl

@@ -73,7 +73,7 @@ function ∇upsample_nearest(x::AbstractArray{T,N}, scales::NTuple{S, <:Integer}
     reshape(mid, outsize)
 end
 
-function ChainRulesCore.rrule(::typeof(upsample_nearest), x::AbstractArray, s::Tuple)
+function rrule(::typeof(upsample_nearest), x::AbstractArray, s::Tuple)
     Ω = upsample_nearest(x, s)
     upsample_nearest_pullback(Δ) = (NO_FIELDS, ∇upsample_nearest(Δ, s), DoesNotExist())
     return Ω, upsample_nearest_pullback
@@ -249,7 +249,7 @@ function ∇upsample_bilinear_whcn!(dx::AbstractArray{T,4}, Δ::AbstractArray{T,
     return dx
 end
 
-function ChainRulesCore.rrule(::typeof(upsample_bilinear), x; size)
+function rrule(::typeof(upsample_bilinear), x; size)
     Ω = upsample_bilinear(x; size=size)
     function upsample_bilinear_pullback(Δ)
         (NO_FIELDS, ∇upsample_bilinear(Δ; size=(Base.size(x,1),Base.size(x,2))))

test/test_utils.jl

@@ -8,7 +8,7 @@ gradtest(f, dims::IntOrTuple...; kw...) =
 Compare numerical gradient and automatic gradient
 given by Zygote. `f` has to be a scalar valued function.
 
-Applies also `ChainRulesTestUtils.rrule_test` if the rrule for `f` is explicitly defined.
+Applies also `ChainRulesTestUtils.test_rrule` if the rrule for `f` is explicitly defined.
 """
 function gradtest(f, xs...; atol=1e-6, rtol=1e-6, fkwargs=NamedTuple(),
                     check_rrule=false,
@@ -19,8 +19,8 @@ function gradtest(f, xs...; atol=1e-6, rtol=1e-6, fkwargs=NamedTuple(),
         y = f(xs...; fkwargs...)
         simil(x) = x isa Number ? randn(rng, typeof(x)) : randn!(rng, similar(x))
         ȳ =  simil(y)
-        xx̄s = [(x, simil(x)) for x in xs]
-        rrule_test(f, ȳ, xx̄s...; fkwargs=fkwargs)
+        xx̄s = [x ⊢ simil(x) for x in xs]
+        test_rrule(f, xx̄s...; fkwargs=fkwargs, output_tangent=ȳ)
     end
 
     if check_broadcast