Fix scaling diagonals for unit triangular matrices

Author: jishnub

src/diagonal.jl

@@ -387,11 +387,10 @@ function _rmul!(A::UpperOrLowerTriangular, D::Diagonal)
         for row in rowstart:rowstop
             P[row, col] *= D.diag[col]
         end
-        if isunit
-            A[col, col] *= D.diag[col]
-        end
     end
-    return A
+    isupper && _setdiag!(P, identity, D.diag)
+    TriWrapper = isupper ? UpperTriangular : LowerTriangular
+    return TriWrapper(P)
 end
 
 function lmul!(D::Diagonal, B::AbstractVecOrMat)
@@ -426,11 +425,10 @@ function _lmul!(D::Diagonal, A::UpperOrLowerTriangular)
         for row in rowstart:rowstop
             P[row, col] = D.diag[col] * P[row, col]
         end
-        if isunit
-            A[col, col] = D.diag[col] * A[col, col]
-        end
     end
-    return A
+    isupper && _setdiag!(P, identity, D.diag)
+    TriWrapper = isupper ? UpperTriangular : LowerTriangular
+    return TriWrapper(P)
 end
 
 @inline function __muldiag_nonzeroalpha!(out, D::Diagonal, B, alpha::Number, beta::Number)