make 2nd order mutation not error

Author: oxinabox

src/extra_rules.jl

@@ -262,3 +262,18 @@ Base.real(z::NoTangent) = z  # TODO should be in CRC, https://github.com/JuliaDi
 
 # Avoid https://github.com/JuliaDiff/ChainRulesCore.jl/pull/495
 ChainRulesCore._backing_error(P::Type{<:Base.Pairs}, G::Type{<:NamedTuple}, E::Type{<:AbstractDict}) = nothing
+
+# Needed for higher order so we don't see the `backing` field of StructuralTangents, just the contents
+# SHould these be in ChainRules/ChainRulesCore?
+# is this always the right behavour, or just because of how we do higher order
+function ChainRulesCore.frule((_, Δ, _, _), ::typeof(getproperty), strct::StructuralTangent, sym::Union{Int,Symbol}, inbounds)
+    return (getproperty(strct, sym, inbounds), getproperty(Δ, sym))
+end
+
+
+function ChainRulesCore.frule((_, ȯbj, _, ẋ), ::typeof(setproperty!), obj::MutableTangent, field, x)
+    ȯbj::MutableTangent
+    y = setproperty!(obj, field, x)
+    ẏ = setproperty!(ȯbj, field, ẋ)
+    return y, ẏ
+end

src/stage1/recurse_fwd.jl

@@ -32,15 +32,15 @@ function (::∂☆new{N})(B::Type, xs::AbstractTangentBundle{N}...) where {N}
     the_primal = _construct(B, primal_args)
     @info "∂☆new{N}"
     the_partials = ntuple(Val{N}()) do ii
-        iith_order_type = ii==1 ? B : Any  # the type of the higher order tangents isn't worth tracking
         tangent_tup = map(x->partial(x, ii), xs)
         tangent = if B<:Tuple
-            Tangent{iith_order_type, typeof(tangent_tup)}(tangent_tup)
+            Tangent{B, typeof(tangent_tup)}(tangent_tup)
         else
-            # TODO support mutation
+            # It is a little dubious using StructuralTangent{B} for >1st order, but it is isomorphic.
+            # Just watch out for order mixing bugs.
             names = fieldnames(B)
             tangent_nt = NamedTuple{names}(tangent_tup)
-            Tangent{iith_order_type, typeof(tangent_nt)}(tangent_nt)
+            StructuralTangent{B}(tangent_nt)
         end
         return tangent
     end

test/forward_mutation.jl

@@ -1,12 +1,26 @@
+# module forward_mutation
 using Diffractor
 using Diffractor: ∂☆, ZeroBundle, TaylorBundle
-using Diffractor: bundle
+using Diffractor: bundle, first_partial, TaylorTangentIndex
+using ChainRulesCore
+using Test
+
 
 mutable struct MDemo1
     x::Float64
 end
 
-∂☆{1}()(ZeroBundle{1}(MDemo1), TaylorBundle{1}(1.5, (1.0,)))
+@testset "construction" begin
+    🍞 = ∂☆{1}()(ZeroBundle{1}(MDemo1), TaylorBundle{1}(1.5, (1.0,)))
+    @test 🍞[TaylorTangentIndex(1)] isa MutableTangent{MDemo1}
+    @test 🍞[TaylorTangentIndex(1)].x == 1.0
+
+    🥯 = ∂☆{2}()(ZeroBundle{2}(MDemo1), TaylorBundle{2}(1.5, (1.0, 10.0)))
+    @test 🥯[TaylorTangentIndex(1)] isa MutableTangent{MDemo1}
+    @test 🥯[TaylorTangentIndex(1)].x == 1.0
+    @test 🥯[TaylorTangentIndex(2)] isa MutableTangent
+    @test 🥯[TaylorTangentIndex(2)].x == 10.0
+end
 
 function double!(val::MDemo1)
     val.x *= 2.0
@@ -16,4 +30,22 @@ function wrap_and_double(x)
     val = MDemo1(x)
     double!(val)
 end
-∂☆{1}()(ZeroBundle{1}(wrap_and_double), TaylorBundle{1}(1.5, (1.0,)))
+🐰 = ∂☆{1}()(ZeroBundle{1}(wrap_and_double), TaylorBundle{1}(1.5, (1.0,)))
+@test first_partial(🐰) isa MutableTangent{MDemo1}
+@test first_partial(🐰).x == 2.0
+
+# second derivative
+🐇 = ∂☆{2}()(ZeroBundle{2}(wrap_and_double), TaylorBundle{2}(1.5, (1.0, 10.0)))
+@test 🐇[TaylorTangentIndex(1)] isa MutableTangent{MDemo1}
+@test 🐇[TaylorTangentIndex(1)].x == 2.0
+@test 🐇[TaylorTangentIndex(2)] isa MutableTangent
+@test 🐇[TaylorTangentIndex(2)] == 0.0  # returns 20
+
+
+
+foo(val) = val^2
+🥖 = ∂☆{2}()(ZeroBundle{2}(foo), TaylorBundle{2}(1.0, (0.0, 10.0)))
+🥖[TaylorTangentIndex(1)] # returns 0
+🥖[TaylorTangentIndex(2)] # returns 20
+
+# end # module