Merge pull request #188 from JuliaDiff/ox/inplaceorder

reorder argument order on InplaceableThunk

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"

docs/Manifest.toml

@@ -11,15 +11,15 @@ uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
 
 [[ChainRulesCore]]
 deps = ["Compat", "LinearAlgebra", "SparseArrays"]
-git-tree-sha1 = "be770c08881f7bb928dfd86d1ba83798f76cf62a"
+git-tree-sha1 = "0b0aa9d61456940511416b59a0e902c57b154956"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "0.10.9"
+version = "0.10.12"
 
 [[ChainRulesTestUtils]]
 deps = ["ChainRulesCore", "Compat", "FiniteDifferences", "LinearAlgebra", "Random", "Test"]
 path = ".."
 uuid = "cdddcdb0-9152-4a09-a978-84456f9df70a"
-version = "0.7.12"
+version = "0.7.13"
 
 [[Compat]]
 deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
@@ -47,19 +47,19 @@ version = "0.8.5"
 
 [[Documenter]]
 deps = ["Base64", "Dates", "DocStringExtensions", "IOCapture", "InteractiveUtils", "JSON", "LibGit2", "Logging", "Markdown", "REPL", "Test", "Unicode"]
-git-tree-sha1 = "621850838b3e74dd6dd047b5432d2e976877104e"
+git-tree-sha1 = "47f13b6305ab195edb73c86815962d84e31b0f48"
 uuid = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
-version = "0.27.2"
+version = "0.27.3"
 
 [[Downloads]]
 deps = ["ArgTools", "LibCURL", "NetworkOptions"]
 uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
 
 [[FiniteDifferences]]
 deps = ["ChainRulesCore", "LinearAlgebra", "Printf", "Random", "Richardson", "StaticArrays"]
-git-tree-sha1 = "bdc9fb1d27a1ccecd2fe8f39c6211524cbe642cb"
+git-tree-sha1 = "12417e4754486a547d98d65293dc0fafdfcc0736"
 uuid = "26cc04aa-876d-5657-8c51-4c34ba976000"
-version = "0.12.13"
+version = "0.12.14"
 
 [[IOCapture]]
 deps = ["Logging", "Random"]
@@ -167,9 +167,9 @@ uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [[StaticArrays]]
 deps = ["LinearAlgebra", "Random", "Statistics"]
-git-tree-sha1 = "745914ebcd610da69f3cb6bf76cb7bb83dcb8c9a"
+git-tree-sha1 = "a43a7b58a6e7dc933b2fa2e0ca653ccf8bb8fd0e"
 uuid = "90137ffa-7385-5640-81b9-e52037218182"
-version = "1.2.4"
+version = "1.2.6"
 
 [[Statistics]]
 deps = ["LinearAlgebra", "SparseArrays"]

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