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, -1),
+            diag(∂y),
+            diag(∂y, 1),
+        )
+    end
+    return y, ∇Tridiagonal
+end

test/rulesets/LinearAlgebra/structured.jl

@@ -161,4 +161,9 @@
             end
         end
     end
+
+    @testset "Tridiagonal" begin
+        test_rrule(Tridiagonal, [1.0, 4.0], [2.0, 3.0, 4.0], [5.0, 3.0])
+        @test pb(10 * res) == (NoTangent(), [10, 40], [20, 30, 40], [50, 30])
+    end
 end