Fix #38 - UnApply for Vector

Author: Keno

src/extra_rules.jl

@@ -256,3 +256,13 @@ end
 @ChainRules.non_differentiable Base.throw(err)
 @ChainRules.non_differentiable Core.Compiler.return_type(args...)
 ChainRulesCore.canonicalize(::NoTangent) = NoTangent()
+
+# Disable thunking at higher order (TODO: These should go into ChainRulesCore)
+function ChainRulesCore.rrule(::Type{Thunk}, thnk)
+    z, ∂z = ∂⃖¹(thnk)
+    z, Δ->(NoTangent(), ∂z(Δ)...)
+end
+
+function ChainRulesCore.rrule(::Type{InplaceableThunk}, add!!, val)
+    val, Δ->(NoTangent(), NoTangent(), Δ)
+end

src/stage1/generated.jl

@@ -336,12 +336,16 @@ function (::∂⃖{N})(::typeof(Core.tuple), args::Vararg{Any, M}) where {N, M}
         )
 end
 
-struct UnApply{Spec}; end
-@generated function (::UnApply{Spec})(Δ) where Spec
+struct UnApply{Spec, Types}; end
+@generated function (::UnApply{Spec, Types})(Δ) where {Spec, Types}
     args = Any[NoTangent(), NoTangent(), :(Δ[1])]
     start = 2
-    for l in Spec
-        push!(args, :(Δ[$(start:(start+l-1))]))
+    for (l, T) in zip(Spec, Types.parameters)
+        if T <: Array
+            push!(args, :([Δ[$(start:(start+l-1))]...]))
+        else
+            push!(args, :(Δ[$(start:(start+l-1))]))
+        end
         start += l
     end
     :(Core.tuple($(args...)))
@@ -362,10 +366,10 @@ end
     a.u(r)
 end
 
-function (this::∂⃖{N})(::typeof(Core._apply_iterate), iterate, f, args::Union{Tuple, NamedTuple}...) where {N}
+function (this::∂⃖{N})(::typeof(Core._apply_iterate), iterate, f, args::Union{Tuple, Vector, NamedTuple}...) where {N}
     @assert iterate === Base.iterate
     x, ∂⃖f = Core._apply_iterate(iterate, this, (f,), args...)
-    return x, ApplyOdd{1, c_order(N)}(UnApply{map(length, args)}(), ∂⃖f)
+    return x, ApplyOdd{1, c_order(N)}(UnApply{map(length, args), typeof(args)}(), ∂⃖f)
 end
 
 

src/stage1/recurse.jl

@@ -571,7 +571,7 @@ function transform!(ci, meth, nargs, sparams, N)
                 if is_accumable(stmt.args[2])
                     accum!(stmt.args[2], nt)
                 end
-            elseif isa(stmt, GlobalRef) || isexpr(stmt, :static_parameter) || isexpr(stmt, :throw_undef_if_not)
+            elseif isa(stmt, GlobalRef) || isexpr(stmt, :static_parameter) || isexpr(stmt, :throw_undef_if_not) || isexpr(stmt, :loopinfo)
                 # We drop gradients for globals and static parameters
             elseif isexpr(stmt, :inbounds)
                 # Nothing to do

test/runtests.jl

@@ -188,4 +188,8 @@ let var"'" = bwd
     @test (x->x^5)'''(1.0) == 60.
 end
 
+# Issue #38 - Splatting arrays
+@test gradient(x -> max(x...), (1,2,3))[1] == (0.0, 0.0, 1.0)
+@test gradient(x -> max(x...), [1,2,3])[1] == [0.0, 0.0, 1.0]
+
 include("pinn.jl")