Merge pull request #758 from JuliaDiff/ChrisRackauckas-patch-1

Add Tridiagonal construction rule
JuliaDiff · Jun 10, 2024 · 78f00cb · 78f00cb
2 parents be9c221 + ce288b9
commit 78f00cb
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/src/rulesets/LinearAlgebra/structured.jl b/src/rulesets/LinearAlgebra/structured.jl
@@ -267,3 +267,16 @@ function rrule(::typeof(logdet), X::Union{Diagonal, AbstractTriangular})
  end
  return y, logdet_pullback
 end
+
+#####
+##### Tridiagonal
+#####
+
+function rrule(::Type{Tridiagonal}, dl, d, du)
+ y = Tridiagonal(dl, d, du)
+ function Tridiagonal_pullback(ȳ)
+ ∂y = unthunk(ȳ)
+ return (NoTangent(), diag(∂y, -1), diag(∂y), diag(∂y, 1))
+ end
+ return y, Tridiagonal_pullback
+end
diff --git a/test/rulesets/LinearAlgebra/structured.jl b/test/rulesets/LinearAlgebra/structured.jl
@@ -161,4 +161,8 @@
  end
  end
  end
+
+ @testset "Tridiagonal" begin
+ test_rrule(Tridiagonal, [1.0, 4.0], [2.0, 3.0, 4.0], [5.0, 3.0])
+ end
 end