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

Add Tridiagonal construction rule

Author: oxinabox

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 = unthunk(ȳ)
+        return (NoTangent(), diag(∂y, -1), diag(∂y), diag(∂y, 1))
+    end
+    return y, Tridiagonal_pullback
+end

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