Add rrules for triu/tril (#201)

* Add rrules for triu/tril

* Add terminal newline

* Increment version number

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "0.6.2"
+version = "0.6.3"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/rulesets/LinearAlgebra/structured.jl

@@ -158,3 +158,29 @@ function rrule(::Type{<:LowerTriangular}, A::AbstractMatrix)
     end
     return LowerTriangular(A), LowerTriangular_pullback
 end
+
+function rrule(::typeof(triu), A::AbstractMatrix, k::Integer)
+    function triu_pullback(ȳ)
+        return (NO_FIELDS, @thunk(triu(ȳ, k)), DoesNotExist())
+    end
+    return triu(A, k), triu_pullback
+end
+function rrule(::typeof(triu), A::AbstractMatrix)
+    function triu_pullback(ȳ)
+        return (NO_FIELDS, @thunk triu(ȳ))
+    end
+    return triu(A), triu_pullback
+end
+
+function rrule(::typeof(tril), A::AbstractMatrix, k::Integer)
+    function tril_pullback(ȳ)
+        return (NO_FIELDS, @thunk(tril(ȳ, k)), DoesNotExist())
+    end
+    return tril(A, k), tril_pullback
+end
+function rrule(::typeof(tril), A::AbstractMatrix)
+    function tril_pullback(ȳ)
+        return (NO_FIELDS, @thunk tril(ȳ))
+    end
+    return tril(A), tril_pullback
+end

test/rulesets/LinearAlgebra/structured.jl

@@ -92,4 +92,11 @@
         n = 5
         rrule_test(T, T(randn(n, n)), (randn(n, n), randn(n, n)))
     end
+    @testset "$Op" for Op in (triu, tril)
+        n = 7
+        rrule_test(Op, randn(n, n), (randn(n, n), randn(n, n)))
+        @testset "k=$k" for k in -2:2
+            rrule_test(Op, randn(n, n), (randn(n, n), randn(n, n)), (k, nothing))
+        end
+    end
 end