Fix codegen of fwd varargs (#137)

* Fix Composites of Structs

* Remove misleading comment

* Fix varargs codegen

Fixes the varargs issue seen in #134, but doesn't quite make nesting
work yet.

* Marked fixed tests as passing

* Fix test

---------

Co-authored-by: Frames White <[email protected]>

Author: Keno

src/codegen/forward.jl

@@ -63,9 +63,8 @@ function fwd_transform!(ci, mi, nargs, N)
     end
 
     meth = mi.def::Method
-    nargs = Int(meth.nargs)
-    for i = 1:nargs
-        if meth.isva && i == nargs
+    for i = 1:meth.nargs
+        if meth.isva && i == meth.nargs
             args = map(i:(nargs+1)) do j::Int
                 emit!(Expr(:call, getfield, SlotNumber(2), j))
             end

src/stage1/forward.jl

@@ -4,7 +4,16 @@ partial(x::TaylorTangent, i) = getfield(getfield(x, :coeffs), i)
 partial(x::UniformTangent, i) = getfield(x, :val)
 partial(x::ProductTangent, i) = ProductTangent(map(x->partial(x, i), getfield(x, :factors)))
 partial(x::AbstractZero, i) = x
-partial(x::CompositeBundle{N, B}, i) where {N, B} = Tangent{B}(map(x->partial(x, i), getfield(x, :tup))...)
+partial(x::CompositeBundle{N, B}, i) where {N, B<:Tuple} = Tangent{B}(map(x->partial(x, i), getfield(x, :tup))...)
+function partial(x::CompositeBundle{N, B}, i) where {N, B}
+    # This is tangent for a struct, but fields partials are each stored in a plain tuple
+    # so we add the names back using the primal `B`
+    # TODO: If required this can be done as a `@generated` function so it is type-stable
+    backing = NamedTuple{fieldnames(B)}(map(x->partial(x, i), getfield(x, :tup)))
+    return Tangent{B, typeof(backing)}(backing)
+end
+
+
 primal(x::AbstractTangentBundle) = x.primal
 primal(z::ZeroTangent) = ZeroTangent()
 

src/stage1/recurse_fwd.jl

@@ -35,7 +35,6 @@ function perform_fwd_transform(world::UInt, source::LineNumberNode,
             Core.svec(:ff, :args), Core.svec(), :(∂☆passthrough(args)))
     end
 
-    # Check if we have an rrule for this function
     sig = Tuple{map(π, args)...}
     mthds = Base._methods_by_ftype(sig, -1, world)
     if mthds === nothing || length(mthds) != 1

test/runtests.jl

@@ -154,7 +154,7 @@ function sin_twice_fwd(x)
     end
 end
 let var"'" = Diffractor.PrimeDerivativeFwd
-    @test_broken sin_twice_fwd'(1.0) == sin'''(1.0)
+    @test sin_twice_fwd'(1.0) == sin'''(1.0)
 end
 
 # Regression tests
@@ -228,7 +228,7 @@ z45, delta45 = frule_via_ad(DiffractorRuleConfig(), (0,1), x -> log(exp(x)), 2)
     @test gradient(x -> sum(sqrt.(atan.(x, transpose(x)))), [1,2,3])[1] ≈ [0.2338, -0.0177, -0.0661] atol=1e-3
     @test gradient(x -> sum(exp.(log.(x))), [1,2,3]) == ([1,1,1],)
 
-    @test_broken gradient(x -> sum((exp∘log).(x)), [1,2,3]) == ([1,1,1],)  # frule_via_ad
+    @test gradient(x -> sum((exp∘log).(x)), [1,2,3]) == ([1,1,1],)  # frule_via_ad
     exp_log(x) = exp(log(x))
     @test gradient(x -> sum(exp_log.(x)), [1,2,3]) == ([1,1,1],)
     @test gradient((x,y) -> sum(x ./ y), [1 2; 3 4], [1,2]) == ([1 1; 0.5 0.5], [-3, -1.75])

test/stage2_fwd.jl

@@ -1,5 +1,6 @@
 module stage2_fwd
     using Diffractor, Test, ChainRulesCore
+
     mysin(x) = sin(x)
     let sin′ = Diffractor.dontuse_nth_order_forward_stage2(Tuple{typeof(mysin), Float64})
         @test sin′(1.0) == cos(1.0)
@@ -17,4 +18,29 @@ module stage2_fwd
         @test isa(self_minus′′, Core.OpaqueClosure{Tuple{Float64}, Float64})
         @test self_minus′′(1.0) == 0.
     end
+
+    @testset "structs" begin
+        struct Foo
+            x
+            y
+        end
+        foo_dub(x) = Foo(x, 2x)
+        dz = Diffractor.∂☆{1}()(Diffractor.ZeroBundle{1}(foo_dub), Diffractor.TaylorBundle{1}(10.0, (π,)))
+        @test Diffractor.first_partial(dz) == Tangent{Foo}(;x=π, y=2π)
+    end
+
+    @testset "mix of vararg and positional args" begin
+        cc(a, x::Vararg) = nothing
+        Diffractor.∂☆{1}()(Diffractor.ZeroBundle{1}(cc), Diffractor.TaylorBundle{1}(10f0, (10.0,)), Diffractor.TaylorBundle{1}(10f0, (10.0,)))
+
+        gg(a, xs...) = nothing
+        Diffractor.∂☆{1}()(Diffractor.ZeroBundle{1}(gg), Diffractor.TaylorBundle{1}(10f0, (1.2,)), Diffractor.TaylorBundle{1}(20f0, (1.1,)))
+    end
+
+
+    @testset "nontrivial nested" begin
+        f(x) = 3x^2
+        g(x) = Diffractor.∂☆{1}()(Diffractor.ZeroBundle{1}(f), Diffractor.TaylorBundle{1}(x, (1.0,)))
+        Diffractor.∂☆{1}()(Diffractor.ZeroBundle{1}(g), Diffractor.TaylorBundle{1}(10f0, (1.0,)))
+    end
 end