From 80b7b2b41a1d6e594dafd62f00b4842f5f87334c Mon Sep 17 00:00:00 2001
From: Oscar Smith <oscar.smith@juliacomputing.com>
Date: Wed, 25 Sep 2024 15:18:43 -0400
Subject: [PATCH] fix

---
 src/internal/linear_solve.jl | 20 +++++++-------------
 1 file changed, 7 insertions(+), 13 deletions(-)

diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl
index 0ad81299e..3e7a191ae 100644
--- a/src/internal/linear_solve.jl
+++ b/src/internal/linear_solve.jl
@@ -80,14 +80,13 @@ function LinearSolverCache(alg, linsolve, A, b, u; stats, kwargs...)
     @bb u_ = copy(u_fixed)
     linprob = LinearProblem(A, b; u0 = u_, kwargs...)
 
-    weight = __init_ones(u_fixed)
     if __hasfield(alg, Val(:precs))
         precs = alg.precs
         Pl_, Pr_ = precs(A, nothing, u, ntuple(Returns(nothing), 6)...)
     else
         precs, Pl_, Pr_ = nothing, nothing, nothing
     end
-    Pl, Pr = __wrapprecs(Pl_, Pr_, weight)
+    Pl, Pr = __wrapprecs(Pl_, Pr_, u)
 
     # Unalias here, we will later use these as caches
     lincache = init(linprob, linsolve; alias_A = false, alias_b = false, Pl, Pr)
@@ -128,10 +127,8 @@ function (cache::LinearSolverCache)(;
     b !== nothing && (cache.lincache.b = b)
     linu !== nothing && __set_lincache_u!(cache, linu)
 
-    Plprev = cache.lincache.Pl isa ComposePreconditioner ? cache.lincache.Pl.outer :
-             cache.lincache.Pl
-    Prprev = cache.lincache.Pr isa ComposePreconditioner ? cache.lincache.Pr.outer :
-             cache.lincache.Pr
+    Plprev = cache.lincache.Pl
+    Prprev = cache.lincache.Pr
 
     if cache.precs === nothing
         _Pl, _Pr = nothing, nothing
@@ -141,10 +138,7 @@ function (cache::LinearSolverCache)(;
     end
 
     if (_Pl !== nothing || _Pr !== nothing)
-        _weight = weight === nothing ?
-                  (cache.lincache.Pr isa Diagonal ? cache.lincache.Pr.diag :
-                   cache.lincache.Pr.inner.diag) : weight
-        Pl, Pr = __wrapprecs(_Pl, _Pr, _weight)
+        Pl, Pr = __wrapprecs(_Pl, _Pr, linu)
         cache.lincache.Pl = Pl
         cache.lincache.Pr = Pr
     end
@@ -242,9 +236,9 @@ function __set_lincache_A(lincache, new_A)
     end
 end
 
-function __wrapprecs(_Pl, _Pr, weight)
-    Pl = _Pl !== nothing ?= _Pl : IdentityOperator(length(weight))
-    Pr = _Pr !== nothing ? _Pr : IdentityOperator(length(weight))
+function __wrapprecs(_Pl, _Pr, u)
+    Pl = _Pl !== nothing ?= _Pl : IdentityOperator(length(u))
+    Pr = _Pr !== nothing ? _Pr : IdentityOperator(length(u))
     return Pl, Pr
 end