Support Enzyme through EnzymeRules

Project.toml

@@ -11,11 +11,13 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 ImplicitDifferentiationChainRulesCoreExt = "ChainRulesCore"
+ImplicitDifferentiationEnzymeExt = "Enzyme"
 ImplicitDifferentiationForwardDiffExt = "ForwardDiff"
 ImplicitDifferentiationZygoteExt = "Zygote"
 
@@ -27,6 +29,7 @@ ChainRulesTestUtils = "1.13.0"
 ComponentArrays = "0.15.27"
 DifferentiationInterface = "0.6.1,0.7"
 Documenter = "1.12.0"
+EnzymeCore = "0.8"
 ExplicitImports = "1"
 FiniteDiff = "2.27.0"
 ForwardDiff = "0.10.36, 1"
@@ -54,6 +57,7 @@ ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"

ext/ImplicitDifferentiationEnzyme.jl

@@ -0,0 +1,80 @@
+module ImplicitDifferentiationForwardDiffExt
+
+using ADTypes: AutoEnzyme
+using EnzymeCore
+import EnzymeCore: EnzymeRules
+using ImplicitDifferentiation:
+    ImplicitFunction,
+    ImplicitFunctionPreparation,
+    IterativeLeastSquaresSolver,
+    build_A,
+    build_Aᵀ,
+    build_B
+
+function EnzymeRules.forward(config, implicit::Const{<:ImplicitFunction}, RT::Type, x, args...)
+	prep = ImplicitFunctionPreparation(eltype(x.val))
+	EnzymeRules.forward(config, implicit, RT, Const(prep), x, args...)
+end
+
+@inline function EnzymeRules.forward(config, implicit::Const{<:ImplicitFunction}, RT::Type, prep::Const{<:ImplicitFunctionPreparation{R}}, x, args...) where R
+	implicit = implicit.val
+	prep = prep.val
+
+	dx = x.dval
+	# dargs = ntuple(length(args)) do i
+		# args[i].dval
+	# end
+
+	x = x.val
+	args = ntuple(length(args)) do i
+        @assert args[i] isa Const
+		args[i].val
+	end
+
+	(; conditions, linear_solver) = implicit
+
+	y, z = implicit(x, args...)
+	c = conditions(x, y, z, args...)
+
+	y0 = zero(y)
+	forward_backend = AutoEnzyme(mode=Forward)
+	reverse_backend = AutoEnzyme(mode=Reverse)
+
+	A = build_A(implicit, prep, x, y, z, c, args...; suggested_backend=forward_backend)
+    B = build_B(implicit, prep, x, y, z, c, args...; suggested_backend=forward_backend)
+    Aᵀ = if linear_solver isa IterativeLeastSquaresSolver
+        build_Aᵀ(implicit, prep, x, y, z, c, args...; suggested_backend=reverse_backend)
+    else
+        nothing
+    end
+
+	if EnzymeRules.width(config) == 1
+		dc = B(dx)
+		dy = linear_solver(A, Aᵀ, dc, y0)::typeof(y0)
+		dz = Enzyme.make_zero(z)
+
+		if EnzymeRules.needs_primal(config)
+			return Duplicated((y, z), (dy, dz)) 
+		else
+			return dy, dz
+		end
+	else
+		dc = map(B, dx)
+		dy = map(dc) do dₖc
+			linear_solver(A, Aᵀ, -dₖc, y0)::typeof(y0)
+		end
+
+		df = ntuple(Val(EnzymeRules.width(config))) do i
+			(dy[i]::typeof(y0), Enzyme.make_zero(z))
+		end
+
+		if EnzymeRules.needs_primal(config)
+			return BatchDuplicated((y, z), df)
+		else
+			return df
+		end
+	end
+end
+
+
+end