Update src/rulesets/LinearAlgebra/structured.jl

JuliaDiff · Nov 23, 2023 · ee952df · ee952df
1 parent 9ba86d8
commit ee952df
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/rulesets/LinearAlgebra/structured.jl b/src/rulesets/LinearAlgebra/structured.jl
@@ -278,9 +278,9 @@ function rrule(::Type{Tridiagonal}, dl, d, du)
     @views function ∇Tridiagonal(∂y)
         return (
             NoTangent(),
-            diag(∂y[2:end, 1:(end - 1)]),
+            diag(∂y, -1),
             diag(∂y),
-            diag(∂y[1:(end - 1), 2:end]),
+            diag(∂y, 1),
         )
     end
     return y, ∇Tridiagonal