Add Tridiagonal construction rule

src/rulesets/LinearAlgebra/structured.jl

@@ -268,3 +268,20 @@ 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)
+    @views function ∇Tridiagonal(∂y)
+        return (
+            NoTangent(),
+            diag(∂y[2:end, 1:(end - 1)]),
+            diag(∂y),
+            diag(∂y[1:(end - 1), 2:end]),
+        )
+    end
+    return y, ∇Tridiagonal
+end

test/rulesets/LinearAlgebra/structured.jl

@@ -161,4 +161,9 @@
             end
         end
     end
+
+    @testset "Tridiagonal" begin
+        res, pb = rrule(Tridiagonal, [1, 4], [2, 3, 4], [5, 3])
+        @test pb(10*res) == (NoTangent(), [10, 40], [20, 30, 40], [50, 30])
+    end
 end