From dd087fb915b4608d97df818a7dd69893edb21e5a Mon Sep 17 00:00:00 2001
From: willtebbutt <wtebbutt@turing.ac.uk>
Date: Wed, 11 Dec 2024 09:44:09 +0000
Subject: [PATCH] Improve trsm stability

---
 src/rrules/blas.jl | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/rrules/blas.jl b/src/rrules/blas.jl
index 7674a5b55..d6576e4d3 100644
--- a/src/rrules/blas.jl
+++ b/src/rrules/blas.jl
@@ -869,7 +869,10 @@ function generate_derived_rrule!!_test_cases(rng_ctor, ::Val{:blas})
         ) do (side, ul, tA, dA, M, N, P)
             t = tA == 'N'
             R = side == 'L' ? M : N
-            As = blas_matrices(rng, P, R, R)
+            As = map(blas_matrices(rng, P, R, R)) do A
+                A[diagind(A)] .+= 1
+                return A
+            end
             Bs = blas_matrices(rng, P, M, N)
             return map_prod(As, Bs) do (A, B)
                 (false, :none, nothing, BLAS.trsm!, side, ul, tA, dA, randn(P), A, B)