Move some code around (NFC)

* Move the BLAS definitions into the `linalg` directory.
* Move some general optimizations into a specific utilities file for
  linear algebra definitions.
* Consolidate structured matrix operations, e.g. diagonal and symmetric,
  into one file.

Author: ararslan

src/ChainRules.jl

@@ -13,11 +13,11 @@ include("rules.jl")
 include("rules/base.jl")
 include("rules/array.jl")
 include("rules/broadcast.jl")
+include("rules/linalg/utils.jl")
+include("rules/linalg/blas.jl")
 include("rules/linalg/dense.jl")
-include("rules/linalg/diagonal.jl")
-include("rules/linalg/symmetric.jl")
+include("rules/linalg/structured.jl")
 include("rules/linalg/factorization.jl")
-include("rules/blas.jl")
 include("rules/nanmath.jl")
 include("rules/specialfunctions.jl")
 

src/rules/blas.jl renamed to src/rules/linalg/blas.jl

@@ -59,33 +59,6 @@ function svd_rev(USV::SVD, Ū::AbstractMatrix, s̄::AbstractVector, V̄::Abstra
     return Ā
 end
 
-function _mulsubtrans!(X::AbstractMatrix{T}, F::AbstractMatrix{T}) where T<:Real
-    k = size(X, 1)
-    @inbounds for j = 1:k, i = 1:j  # Iterate the upper triangle
-        if i == j
-            X[i,i] = zero(T)
-        else
-            X[i,j], X[j,i] = F[i,j] * (X[i,j] - X[j,i]), F[j,i] * (X[j,i] - X[i,j])
-        end
-    end
-    X
-end
-
-function _eyesubx!(X::AbstractMatrix{T}) where T<:Real
-    n, m = size(X)
-    @inbounds for j = 1:m, i = 1:n
-        X[i,j] = (i == j) - X[i,j]
-    end
-    X
-end
-
-function _add!(X::AbstractMatrix{T}, Y::AbstractMatrix{T}) where T<:Real
-    @inbounds for i = eachindex(X, Y)
-        X[i] += Y[i]
-    end
-    X
-end
-
 #####
 ##### `cholesky`
 #####

src/rules/linalg/diagonal.jl

@@ -0,0 +1,18 @@
+# Structured matrices
+
+#####
+##### `Diagonal`
+#####
+
+rrule(::Type{<:Diagonal}, d::AbstractVector) = Diagonal(d), Rule(diag)
+
+rrule(::typeof(diag), A::AbstractMatrix) = diag(A), Rule(Diagonal)
+
+#####
+##### `Symmetric`
+#####
+
+rrule(::Type{<:Symmetric}, A::AbstractMatrix) = Symmetric(A), Rule(_symmetric_back)
+
+_symmetric_back(ΔΩ) = UpperTriangular(ΔΩ) + LowerTriangular(ΔΩ)' - Diagonal(ΔΩ)
+_symmetric_back(ΔΩ::Union{Diagonal,UpperTriangular}) = ΔΩ

src/rules/linalg/factorization.jl

@@ -0,0 +1,32 @@
+# Some utility functions for optimizing linear algebra operations that aren't specific
+# to any particular rule definition
+
+# F .* (X - X'), overwrites X
+function _mulsubtrans!(X::AbstractMatrix{T}, F::AbstractMatrix{T}) where T<:Real
+    k = size(X, 1)
+    @inbounds for j = 1:k, i = 1:j  # Iterate the upper triangle
+        if i == j
+            X[i,i] = zero(T)
+        else
+            X[i,j], X[j,i] = F[i,j] * (X[i,j] - X[j,i]), F[j,i] * (X[j,i] - X[i,j])
+        end
+    end
+    X
+end
+
+# I - X, overwrites X
+function _eyesubx!(X::AbstractMatrix)
+    n, m = size(X)
+    @inbounds for j = 1:m, i = 1:n
+        X[i,j] = (i == j) - X[i,j]
+    end
+    X
+end
+
+# X + Y, overwrites X
+function _add!(X::AbstractVecOrMat{T}, Y::AbstractVecOrMat{T}) where T<:Real
+    @inbounds for i = eachindex(X, Y)
+        X[i] += Y[i]
+    end
+    X
+end

src/rules/linalg/structured.jl

@@ -1,4 +1,4 @@
-@testset "diagonal" begin
+@testset "Structured Matrices" begin
     @testset "Diagonal" begin
         rng, N = MersenneTwister(123456), 3
         rrule_test(Diagonal, randn(rng, N, N), (randn(rng, N), randn(rng, N)))
@@ -14,4 +14,8 @@
         rrule_test(diag, randn(rng, N), (randn(rng, N, N), Diagonal(randn(rng, N))))
         rrule_test(diag, randn(rng, N), (Diagonal(randn(rng, N)), Diagonal(randn(rng, N))))
     end
+    @testset "Symmetric" begin
+        rng, N = MersenneTwister(123456), 3
+        rrule_test(Symmetric, randn(rng, N, N), (randn(rng, N, N), randn(rng, N, N)))
+    end
 end

src/rules/linalg/symmetric.jl

@@ -17,8 +17,7 @@ include("test_util.jl")
         include(joinpath("rules", "array.jl"))
         @testset "linalg" begin
             include(joinpath("rules", "linalg", "dense.jl"))
-            include(joinpath("rules", "linalg", "diagonal.jl"))
-            include(joinpath("rules", "linalg", "symmetric.jl"))
+            include(joinpath("rules", "linalg", "structured.jl"))
             include(joinpath("rules", "linalg", "factorization.jl"))
         end
         include(joinpath("rules", "broadcast.jl"))