Performance tweaks (#47)

JuliaDecisionFocusedLearning · May 24, 2023 · 3f3bc10 · 3f3bc10
1 parent a710a81
commit 3f3bc10
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 3 deletions.
diff --git a/ext/ImplicitDifferentiationChainRulesExt.jl b/ext/ImplicitDifferentiationChainRulesExt.jl
@@ -51,7 +51,8 @@ function (implicit_pullback::ImplicitPullback)((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_vec = Bᵀ_op * dF_vec
+    dx_vec .*= -1
     dx = reshape(dx_vec, size(x))
     return (NoTangent(), dx)
 end

diff --git a/ext/ImplicitDifferentiationForwardDiffExt.jl b/ext/ImplicitDifferentiationForwardDiffExt.jl
@@ -33,7 +33,8 @@ function (implicit::ImplicitFunction)(
 
     dy = map(1:N) do k
         dₖx_vec = vec(partials.(x_and_dx, k))
-        dₖy_vec, stats = linear_solver(A_op, -(B_op * dₖx_vec))
+        dₖy_vec, stats = linear_solver(A_op, B_op * dₖx_vec)
+        dₖy_vec .*= -1
         check_solution(linear_solver, stats)
         reshape(dₖy_vec, size(y))
     end

diff --git a/src/ImplicitDifferentiation.jl b/src/ImplicitDifferentiation.jl
@@ -2,7 +2,7 @@ module ImplicitDifferentiation
 
 using AbstractDifferentiation: LazyJacobian, ReverseRuleConfigBackend, lazy_jacobian
 using Krylov: KrylovStats, gmres
-using LinearOperators: LinearOperator
+using LinearOperators: LinearOperators, LinearOperator
 using Requires: @require
 
 include("utils.jl")

diff --git a/src/utils.jl b/src/utils.jl
@@ -57,3 +57,8 @@ function (pbm::PullbackMul!)(res::AbstractVector, δoutput_vec::AbstractVector)
     δinput = only(pbm.pullback(δoutput))
     return res .= vec(δinput)
 end
+
+## Override this function from LinearOperators to avoid generating the whole methods table
+
+LinearOperators.get_nargs(pfm::PushforwardMul!) = 1
+LinearOperators.get_nargs(pbm::PullbackMul!) = 1