Improve type stability (#41)

* Improve type stability

* Update ImplicitDifferentiationForwardDiffExt.jl

* Add type inference tests and bump version to 0.4.1

---------

Co-authored-by: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>

Author: YingboMa

CITATION.bib

@@ -2,7 +2,7 @@ @misc{ImplicitDifferentiation.jl
 	author  = {Guillaume Dalle, Mohamed Tarek and contributors},
 	title   = {ImplicitDifferentiation.jl},
 	url     = {https://github.com/gdalle/ImplicitDifferentiation.jl},
-	version = {v0.4.0},
+	version = {v0.4.1},
 	year    = {2023},
-	month   = {4}
+	month   = {5}
 }

Project.toml

@@ -1,7 +1,7 @@
 name = "ImplicitDifferentiation"
 uuid = "57b37032-215b-411a-8a7c-41a003a55207"
 authors = ["Guillaume Dalle", "Mohamed Tarek and contributors"]
-version = "0.4.0"
+version = "0.4.1"
 
 [deps]
 AbstractDifferentiation = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d"
@@ -10,6 +10,14 @@ Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
 LinearOperators = "5c8ed15e-5a4c-59e4-a42b-c7e8811fb125"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 
+[weakdeps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+
+[extensions]
+ImplicitDifferentiationChainRulesExt = "ChainRulesCore"
+ImplicitDifferentiationForwardDiffExt = "ForwardDiff"
+
 [compat]
 AbstractDifferentiation = "0.5"
 ChainRulesCore = "1.14"
@@ -19,10 +27,6 @@ LinearOperators = "2.2"
 Requires = "1.3"
 julia = "1.6"
 
-[extensions]
-ImplicitDifferentiationChainRulesExt = "ChainRulesCore"
-ImplicitDifferentiationForwardDiffExt = "ForwardDiff"
-
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -43,7 +47,3 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
 test = ["Aqua", "ChainRulesCore", "ChainRulesTestUtils", "Documenter", "ForwardDiff", "JET", "JuliaFormatter", "LinearAlgebra", "NLsolve", "Optim", "Pkg", "Random", "SparseArrays", "Test", "Zygote"]
-
-[weakdeps]
-ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"

examples/0_basic.jl

@@ -4,8 +4,11 @@
 In this example, we demonstrate the basics of our package on a simple function that is not amenable to automatic differentiation.
 =#
 
+using ChainRulesCore  #src
+using ChainRulesTestUtils  #src
 using ForwardDiff
 using ImplicitDifferentiation
+using JET  #src
 using LinearAlgebra
 using Random
 using Test  #src
@@ -168,8 +171,16 @@ JJ = Diagonal(0.5 ./ sqrt.(vec(X)))  #src
 @test ForwardDiff.jacobian(first ∘ implicit, X) ≈ JJ  #src
 @test Zygote.jacobian(first ∘ implicit, X)[1] ≈ JJ  #src
 
-# Skipped because of https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/232  #src
+# Skipped because of https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/232 and because it detects weird type instabilities  #src
 @testset verbose = true "ChainRulesTestUtils.jl" begin  #src
     @test_skip test_rrule(implicit, x)  #src
     @test_skip test_rrule(implicit, X)  #src
 end  #src
+
+x_and_dx = [ForwardDiff.Dual(x[i], (0, 0)) for i in eachindex(x)]  #src
+@inferred implicit(x_and_dx)  #src
+
+rc = Zygote.ZygoteRuleConfig()  #src
+_, pullback = @inferred rrule(rc, implicit, x)  #src
+dy, dz = zero(implicit(x)[1]), 0
+@inferred pullback((dy, dz))

ext/ImplicitDifferentiationChainRulesExt.jl

@@ -25,19 +25,35 @@ function ChainRulesCore.rrule(
     backend = ReverseRuleConfigBackend(rc)
     pbA = pullback_function(backend, _y -> conditions(x, _y, z; kwargs...), y)
     pbB = pullback_function(backend, _x -> conditions(_x, y, z; kwargs...), x)
-    Aᵀ_op = LinearOperator(R, m, m, false, false, PullbackMul!(pbA, size(y)))
-    Bᵀ_op = LinearOperator(R, n, m, false, false, PullbackMul!(pbB, size(y)))
-
-    function implicit_pullback((dy, dz))
-        dy_vec = convert(Vector{R}, vec(unthunk(dy)))
-        dF_vec, stats = linear_solver(Aᵀ_op, dy_vec)
-        check_solution(linear_solver, stats)
-        dx_vec = -(Bᵀ_op * dF_vec)
-        dx = reshape(dx_vec, size(x))
-        return (NoTangent(), dx)
-    end
+    pbmA = PullbackMul!(pbA, size(y))
+    pbmB = PullbackMul!(pbB, size(y))
+    Aᵀ_op = LinearOperator(R, m, m, false, false, pbmA)
+    Bᵀ_op = LinearOperator(R, n, m, false, false, pbmB)
+    implicit_pullback = ImplicitPullback(Aᵀ_op, Bᵀ_op, linear_solver, x)
 
     return (y, z), implicit_pullback
 end
 
+struct ImplicitPullback{A,B,L,X}
+    Aᵀ_op::A
+    Bᵀ_op::B
+    linear_solver::L
+    x::X
+end
+
+function (implicit_pullback::ImplicitPullback)((dy, dz))
+    Aᵀ_op = implicit_pullback.Aᵀ_op
+    Bᵀ_op = implicit_pullback.Bᵀ_op
+    linear_solver = implicit_pullback.linear_solver
+    x = implicit_pullback.x
+    R = eltype(x)
+
+    dy_vec = convert(Vector{R}, vec(unthunk(dy)))
+    dF_vec, stats = linear_solver(Aᵀ_op, dy_vec)
+    check_solution(linear_solver, stats)
+    dx_vec = -(Bᵀ_op * dF_vec)
+    dx = reshape(dx_vec, size(x))
+    return (NoTangent(), dx)
+end
+
 end

ext/ImplicitDifferentiationForwardDiffExt.jl

@@ -38,11 +38,15 @@ function (implicit::ImplicitFunction)(
         reshape(dₖy_vec, size(y))
     end
 
-    y_and_dy = map(eachindex(y)) do i
-        Dual{T}(y[i], Partials(Tuple(dy[k][i] for k in 1:N)))
+    y_and_dy = let y = y, dy = dy
+        map(eachindex(y)) do i
+            Dual{T}(y[i], Partials(ntuple(k -> dy[k][i], Val(N))))
+        end
     end
 
-    z_and_dz = Dual{T}(z, Partials(Tuple(zero(z) for k in 1:N)))
+    z_and_dz = let z = z
+        Dual{T}(z, Partials(ntuple(_ -> zero(z), Val(N))))
+    end
 
     return y_and_dy, z_and_dz
 end