reorder argument order on InplaceableThink

Author: oxinabox

Project.toml

@@ -11,7 +11,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
-ChainRulesCore = "0.10.4"
+ChainRulesCore = "0.10.12"
 Compat = "3"
 FiniteDifferences = "0.12.12"
 julia = "1"

test/check_result.jl

@@ -17,11 +17,11 @@ end
         check(11.0, ZeroTangent())
         check([10.0, 20.0], @thunk([2.0, 0.0]))
 
-        check(12.0, InplaceableThunk(@thunk(2.0), X̄ -> error("Should not have in-placed")))
+        check(12.0, InplaceableThunk(X̄ -> error("Should not have in-placed"), @thunk(2.0)))
 
-        check([10.0, 20.0], InplaceableThunk(@thunk([2.0, 0.0]), X̄ -> (X̄[1] += 2.0; X̄)))
+        check([10.0, 20.0], InplaceableThunk(X̄ -> (X̄[1] += 2.0; X̄), @thunk([2.0, 0.0])))
         @test fails() do
-            check([10.0, 20.0], InplaceableThunk(@thunk([2.0, 0.0]), X̄ -> (X̄[1] += 3.0; X̄)))
+            check([10.0, 20.0], InplaceableThunk(X̄ -> (X̄[1] += 3.0; X̄), @thunk([2.0, 0.0])))
         end
     end
 

test/testers.jl

@@ -96,14 +96,15 @@ struct MySpecialConfig <: RuleConfig{Union{MySpecialTrait}} end
         @testset "Correct definitions" begin
             local inplace_used
             function ChainRulesCore.frule((_, ẋ), ::typeof(identity), x::Array)
-                ẏ = InplaceableThunk(@thunk(ẋ), ȧ -> (inplace_used = true; ȧ .+= ẋ))
+                ẏ = InplaceableThunk(ȧ -> (inplace_used = true; ȧ .+= ẋ), @thunk(ẋ))
                 return identity(x), ẏ
             end
             function ChainRulesCore.rrule(::typeof(identity), x::Array)
                 function identity_pullback(ȳ)
-                    x̄_ret = InplaceableThunk(
-                        @thunk(ȳ), ā -> (inplace_used = true; ā .+= ȳ)
-                    )
+                    x̄_ret = InplaceableThunk(@thunk(ȳ)) do ā
+                        inplace_used = true
+                        ā .+= ȳ
+                    end
                     return (NoTangent(), x̄_ret)
                 end
                 return identity(x), identity_pullback
@@ -122,14 +123,14 @@ struct MySpecialConfig <: RuleConfig{Union{MySpecialTrait}} end
             my_identity(value) = value  # we will define bad rules on this
             function ChainRulesCore.frule((_, ẋ), ::typeof(my_identity), x::Array)
                 # only the in-place part is incorrect
-                ẏ = InplaceableThunk(@thunk(ẋ), ȧ -> ȧ .+= 200 .* ẋ)
+                ẏ = InplaceableThunk(ȧ -> ȧ .+= 200 .* ẋ, @thunk(ẋ))
                 return my_identity(x), ẏ
             end
             function ChainRulesCore.rrule(::typeof(my_identity), x::Array)
                 x_dims = size(x)
                 function my_identity_pullback(ȳ)
                     # only the in-place part is incorrect
-                    x̄_ret = InplaceableThunk(@thunk(ȳ), ā -> ā .+= 200 .* ȳ)
+                    x̄_ret = InplaceableThunk(ā -> ā .+= 200 .* ȳ, @thunk(ȳ))
                     return (NoTangent(), x̄_ret)
                 end
                 return my_identity(x), my_identity_pullback