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.
JuliaDiff · Jun 11, 2019 · 4d46dc7 · 4d46dc7
1 parent 0e551fb
commit 4d46dc7
Show file tree

Hide file tree

Showing 7 changed files with 53 additions and 36 deletions.
diff --git a/src/ChainRules.jl b/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")
 

diff --git a/src/rules/blas.jl → src/rules/linalg/blas.jl b/src/rules/blas.jl → src/rules/linalg/blas.jl
diff --git a/src/rules/linalg/diagonal.jl b/src/rules/linalg/diagonal.jl
diff --git a/src/rules/linalg/factorization.jl b/src/rules/linalg/factorization.jl
@@ -59,33 +59,6 @@ function svd_rev(USV::SVD, Ū::AbstractMatrix, s̄::AbstractVector, V̄::Abstra
     return Ā
 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`
 #####

diff --git a/src/rules/linalg/structured.jl b/src/rules/linalg/structured.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}) = ΔΩ
diff --git a/src/rules/linalg/symmetric.jl b/src/rules/linalg/symmetric.jl
diff --git a/src/rules/linalg/utils.jl b/src/rules/linalg/utils.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