Fix Backend Compatibility

Project.toml

@@ -2,15 +2,23 @@ name = "DiffTests"
 uuid = "de460e47-3fe3-5279-bb4a-814414816d5d"
 keywords = ["automatic differentiation", "test"]
 license = "MIT"
-version = "0.1.2"
+version = "0.1.3"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
+[weakdeps]
+MKLSparse = "0c723cd3-b8cd-5d40-b370-ba682dde9aae"
+
+[extensions]
+# some operations (ldiv!) are only supported by MKLSparse
+DiffTestsMKLSparseExt = ["MKLSparse"]
+
 [compat]
 julia = "1"
+MKLSparse = "1.2"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

ext/DiffTestsMKLSparseExt.jl

@@ -0,0 +1,16 @@
+module DiffTestsMKLSparseExt
+
+import DiffTests: test_matrix, register
+import DiffTests.LinearAlgebra: UpperTriangular, LowerTriangular
+import DiffTests.SparseArrays: sparse
+
+#sparse_ldiv(x::AbstractVecOrMat) = sparse(test_matrix(x)) \ x # arbitrary sparse matrix not supported
+sp_utriag_ldiv(x::AbstractVecOrMat) = UpperTriangular(sparse(test_matrix(x))) \ x
+sp_ltriag_ldiv(x::AbstractVecOrMat) = LowerTriangular(sparse(test_matrix(x))) \ x
+
+function __init__()
+    register(:VECTOR_TO_VECTOR_FUNCS, [sp_utriag_ldiv, sp_ltriag_ldiv])
+    register(:MATRIX_TO_MATRIX_FUNCS, [sp_utriag_ldiv, sp_ltriag_ldiv])
+end
+
+end

src/DiffTests.jl

@@ -1,9 +1,36 @@
 module DiffTests
 
-using LinearAlgebra: det, norm, dot, tr, Diagonal, LowerTriangular, UpperTriangular
+import LinearAlgebra
+import SparseArrays
+
+using LinearAlgebra: det, norm, dot, tr, lu, Diagonal, LowerTriangular, UpperTriangular
 using SparseArrays: sparse
 using Statistics: mean
 
+# registers given functions in the specified functions registry
+function register(registry::AbstractVector, funcs::AbstractVector)
+    for f in funcs
+        isnothing(findfirst(==(f), registry)) && push!(registry, f)
+    end
+    return registry
+end
+
+function register(registry_id::Symbol, funcs::AbstractVector)
+    registry = nothing
+    try
+        registry = getproperty(@__MODULE__, registry_id)
+    catch e
+        if e isa UndefVarError
+            registry = Vector{Function}()
+            setproperty!(@__MODULE__, registry_id, registry)
+        else
+            rethrow(e)
+        end
+    end
+    register(registry, funcs)
+    return registry
+end
+
 #=
 These functions are organized in sets based on input/output type. They are unary and not
 in-place unless otherwised specified. These functions have been written with the following
@@ -28,12 +55,9 @@ num2num_3(x) = 10.31^(x + x) - x
 num2num_4(x) = 1
 num2num_5(x) = 1. / (1. + exp(-x))
 
-const NUMBER_TO_NUMBER_FUNCS = (num2num_1, num2num_2, num2num_3,
-                                num2num_4, num2num_5, identity)
-
-#######################
-# f(x::Number)::Array #
-#######################
+###############################
+# f(x::Number)::AbstractArray #
+###############################
 
 function num2arr_1(x)
     return reshape([num2num_1(x),
@@ -46,11 +70,9 @@ function num2arr_1(x)
                     num2num_3(x)], 2, 2, 2)
 end
 
-const NUMBER_TO_ARRAY_FUNCS = (num2arr_1,)
-
-####################################
-# f!(y::Array, x::Number)::Nothing #
-####################################
+############################################
+# f!(y::AbstractArray, x::Number)::Nothing #
+############################################
 
 function num2arr_1!(y, x)
     fill!(y, zero(x))
@@ -60,18 +82,16 @@ function num2arr_1!(y, x)
     return nothing
 end
 
-const INPLACE_NUMBER_TO_ARRAY_FUNCS = (num2arr_1!,)
-
-########################
-# f(x::Vector)::Number #
-########################
+################################
+# f(x::AbstractVector)::Number #
+################################
 
 vec2num_1(x) = (exp(x[1]) + log(x[3]) * x[4]) / x[5]
 vec2num_2(x) = x[1]*x[2] + sin(x[1])
 vec2num_3(x) = norm(x' .* x)
 vec2num_4(x) = ((sum(x) + prod(x)); 1)
 vec2num_5(x) = sum((-x).^3)
-vec2num_6(x) = sum([ifelse(i > 0, i, 0) for i in x])
+vec2num_6(x) = sum([ifelse(i > 0, i, zero(eltype(x))) for i in x])
 vec2num_7(x) = sum(map(y -> x[1] * y, x))
 
 function rosenbrock_1(x)
@@ -114,16 +134,11 @@ end
 
 self_weighted_logit(x) = inv(1.0 + exp(-dot(x, x)))
 
-nested_array_mul(x) = sum(sum(x[1] * [[x[2], x[3]]]))
+nested_array_mul(x) = sum(sum(x[1] * [[x[2], x[3]]]))::Base.promote_op(LinearAlgebra.matprod, eltype(x), eltype(x))
 
-const VECTOR_TO_NUMBER_FUNCS = (vec2num_1, vec2num_2,  vec2num_3, vec2num_4, vec2num_5,
-                                vec2num_6, vec2num_7, rosenbrock_1, rosenbrock_2,
-                                rosenbrock_3, rosenbrock_4, ackley, self_weighted_logit,
-                                nested_array_mul, first)
-
-########################
-# f(x::Matrix)::Number #
-########################
+################################
+# f(x::AbstractMatrix)::Number #
+################################
 
 mat2num_1(x) = det(first(x) * inv(x * x) + x)
 
@@ -144,8 +159,6 @@ mat2num_4(x) = mean(sum(sin.(x) * x, dims=2))
 
 softmax(x) = sum(exp.(x) ./ sum(exp.(x), dims=2))
 
-const MATRIX_TO_NUMBER_FUNCS = (det, mat2num_1, mat2num_2, mat2num_3, mat2num_4, softmax)
-
 ####################
 # binary broadcast #
 ####################
@@ -157,28 +170,26 @@ const BINARY_BROADCAST_OPS = ((a, b) -> broadcast(+, a, b),
                               (a, b) -> broadcast(\, a, b),
                               (a, b) -> broadcast(^, a, b))
 
-#################################
-# f(::Matrix, ::Matrix)::Number #
-#################################
+#########################################################
+# f(::AbstractMatrix, ::AbstractMatrix)::AbstractMatrix #
+#########################################################
 
-const BINARY_MATRIX_TO_MATRIX_FUNCS = (+, -, *, /, \,
-                                       BINARY_BROADCAST_OPS...,
-                                       (a, b) -> a * transpose(b), (a, b) -> transpose(a) * b, (a, b) -> transpose(a) * transpose(b),
-                                       (a, b) -> a * adjoint(b), (a, b) -> adjoint(a) * b, (a, b) -> adjoint(a) * adjoint(b))
+# Julia LinearAlgebra does not support matrix\matrix,
+# one needs to compute A factorization first
+ldiv_lu(A::AbstractMatrix, B::AbstractArray) = lu(A) \ B
+rdiv_lu(A::AbstractArray, B::AbstractMatrix) = A / lu(B)
 
-###########################################
-# f(::Matrix, ::Matrix, ::Matrix)::Number #
-###########################################
+###################################################################
+# f(::AbstractMatrix, ::AbstractMatrix, ::AbstractMatrix)::Number #
+###################################################################
 
 relu(x) = log.(1.0 .+ exp.(x))
 sigmoid(n) = 1. / (1. + exp.(-n))
 neural_step(x1, w1, w2) = sigmoid(dot(w2[1:size(w1, 2)], relu(w1 * x1[1:size(w1, 2)])))
 
-const TERNARY_MATRIX_TO_NUMBER_FUNCS = (neural_step,)
-
-###################################
-# f!(y::Array, x::Array)::Nothing #
-###################################
+###################################################
+# f!(y::AbstractArray, x::AbstractArray)::Nothing #
+###################################################
 # Credit for `chebyquad!`, `brown_almost_linear!`, and `trigonometric!` goes to
 # Kristoffer Carlsson (@KristofferC).
 
@@ -247,49 +258,39 @@ function mutation_test_2!(y, x)
     return nothing
 end
 
-const INPLACE_ARRAY_TO_ARRAY_FUNCS = (chebyquad!, brown_almost_linear!, trigonometric!,
-                                      mutation_test_1!, mutation_test_2!)
-
-############################
-# f(x::VecOrMat)::VecOrMat #
-############################
+############################################
+# f(x::AbstractVecOrMat)::AbstractVecOrMat #
+############################################
 
 diag_matrix(::Type{T}, n::Integer) where T<:Real =
     Diagonal(LinRange(convert(T, 0.01), convert(T, 100.0), n))
-diag_matrix(x::VecOrMat) = diag_matrix(Float64, size(x, 1))
+diag_matrix(x::AbstractVecOrMat) = diag_matrix(Float64, size(x, 1))
 
 lehmer_matrix(::Type{T}, n::Integer) where T<:Real =
     [convert(T, min(i, j)/max(i, j)) for i in 1:n, j in 1:n]
-lehmer_matrix(x::VecOrMat) = lehmer_matrix(Float64, size(x, 1))
+lehmer_matrix(x::AbstractVecOrMat)::Matrix{Float64} = lehmer_matrix(Float64, size(x, 1))
 
-test_matrix = lehmer_matrix
+test_matrix(x) = lehmer_matrix(x)
 
 # left multiplication by a constant matrix
-diag_lmul(x::VecOrMat) = diag_matrix(x) * x
-dense_lmul(x::VecOrMat) = test_matrix(x) * x
-utriag_lmul(x::VecOrMat) =  UpperTriangular(test_matrix(x)) * x
-ltriag_lmul(x::VecOrMat) =  LowerTriangular(test_matrix(x)) * x
-sparse_lmul(x::VecOrMat) = sparse(test_matrix(x)) * x
-sp_utriag_lmul(x::VecOrMat) = UpperTriangular(sparse(test_matrix(x))) * x
-sp_ltriag_lmul(x::VecOrMat) = LowerTriangular(sparse(test_matrix(x))) * x
+diag_lmul(x::AbstractVecOrMat) = diag_matrix(x) * x
+dense_lmul(x::AbstractVecOrMat) = test_matrix(x) * x
+utriag_lmul(x::AbstractVecOrMat) = UpperTriangular(test_matrix(x)) * x
+ltriag_lmul(x::AbstractVecOrMat) = LowerTriangular(test_matrix(x)) * x
+
+sparse_lmul(x::AbstractVecOrMat) = sparse(test_matrix(x)) * x
+sp_utriag_lmul(x::AbstractVecOrMat) = UpperTriangular(sparse(test_matrix(x))) * x
+sp_ltriag_lmul(x::AbstractVecOrMat) = LowerTriangular(sparse(test_matrix(x))) * x
 
 # left division by a constant matrix
-diag_ldiv(x::VecOrMat) = diag_matrix(x) \ x
-dense_ldiv(x::VecOrMat) = test_matrix(x) \ x
-utriag_ldiv(x::VecOrMat) =  UpperTriangular(test_matrix(x)) \ x
-ltriag_ldiv(x::VecOrMat) =  LowerTriangular(test_matrix(x)) \ x
-sparse_ldiv(x::VecOrMat) = sparse(test_matrix(x)) \ x
-sp_utriag_ldiv(x::VecOrMat) = UpperTriangular(sparse(test_matrix(x))) \ x
-sp_ltriag_ldiv(x::VecOrMat) = LowerTriangular(sparse(test_matrix(x))) \ x
-
-const VECTOR_TO_VECTOR_FUNCS = (diag_lmul, dense_lmul, utriag_lmul, ltriag_lmul,
-                                sparse_lmul, sp_utriag_lmul, sp_ltriag_lmul,
-                                diag_ldiv, utriag_ldiv, ltriag_ldiv,
-                                sparse_ldiv, sp_utriag_ldiv, sp_ltriag_ldiv,)
-
-######################
-# f(x::Array)::Array #
-######################
+diag_ldiv(x::AbstractVecOrMat) = diag_matrix(x) \ x
+dense_ldiv(x::AbstractVecOrMat) = test_matrix(x) \ x
+utriag_ldiv(x::AbstractVecOrMat) = UpperTriangular(test_matrix(x)) \ x
+ltriag_ldiv(x::AbstractVecOrMat) = LowerTriangular(test_matrix(x)) \ x
+
+######################################
+# f(x::AbstractArray)::AbstractArray #
+######################################
 
 chebyquad(x) = (y = fill(zero(eltype(x)), size(x)); chebyquad!(y, x); return y)
 
@@ -305,17 +306,60 @@ arr2arr_1(x) = (sum(x .* x); fill(zero(eltype(x)), size(x)))
 
 arr2arr_2(x) = x[1, :] .+ x[1, :] .+ first(x)
 
-const ARRAY_TO_ARRAY_FUNCS = (-, chebyquad, brown_almost_linear, trigonometric, arr2arr_1,
-                              arr2arr_2, mutation_test_1, mutation_test_2, identity)
+#######################################
+# f(::AbstractMatrix)::AbstractMatrix #
+#######################################
+
+function __init__()
+    register(:NUMBER_TO_NUMBER_FUNCS,
+             [num2num_1, num2num_2, num2num_3,
+              num2num_4, num2num_5, identity])
+
+    register(:NUMBER_TO_ARRAY_FUNCS, [num2arr_1,])
+
+    register(:INPLACE_NUMBER_TO_ARRAY_FUNCS, [num2arr_1!,])
+
+    register(:VECTOR_TO_NUMBER_FUNCS,
+        [vec2num_1, vec2num_2,  vec2num_3, vec2num_4, vec2num_5,
+         vec2num_6, vec2num_7, rosenbrock_1, rosenbrock_2,
+         rosenbrock_3, rosenbrock_4, ackley, self_weighted_logit,
+         nested_array_mul, first])
 
-#######################
-# f(::Matrix)::Matrix #
-#######################
+    register(:MATRIX_TO_NUMBER_FUNCS,
+        [det, mat2num_1, mat2num_2, mat2num_3, mat2num_4, softmax])
 
-const MATRIX_TO_MATRIX_FUNCS = (inv,
-                                diag_lmul, dense_lmul, utriag_lmul, ltriag_lmul,
-                                sparse_lmul, sp_utriag_lmul, sp_ltriag_lmul,
-                                diag_ldiv, utriag_ldiv, ltriag_ldiv,
-                                sparse_ldiv, sp_utriag_ldiv, sp_ltriag_ldiv,)
+    register(:TERNARY_MATRIX_TO_NUMBER_FUNCS, [neural_step,])
+
+    register(:VECTOR_TO_VECTOR_FUNCS,
+        [diag_lmul, dense_lmul, utriag_lmul, ltriag_lmul,
+         sparse_lmul, sp_utriag_lmul, sp_ltriag_lmul,
+         diag_ldiv,])
+
+    register(:MATRIX_TO_MATRIX_FUNCS,
+        [inv,
+         diag_lmul, dense_lmul, utriag_lmul, ltriag_lmul,
+         sparse_lmul, sp_utriag_lmul, sp_ltriag_lmul,
+         diag_ldiv,])
+
+    register(:ARRAY_TO_ARRAY_FUNCS,
+        [-, chebyquad, brown_almost_linear, trigonometric, arr2arr_1,
+         arr2arr_2, mutation_test_1, mutation_test_2, identity])
+
+    register(:INPLACE_ARRAY_TO_ARRAY_FUNCS,
+        [chebyquad!, brown_almost_linear!, trigonometric!,
+         mutation_test_1!, mutation_test_2!])
+
+    register(:BINARY_MATRIX_TO_MATRIX_FUNCS,
+        [+, -, *, rdiv_lu, ldiv_lu,
+         BINARY_BROADCAST_OPS...,
+         (a, b) -> a * transpose(b), (a, b) -> transpose(a) * b, (a, b) -> transpose(a) * transpose(b),
+         (a, b) -> a * adjoint(b), (a, b) -> adjoint(a) * b, (a, b) -> adjoint(a) * adjoint(b)])
+
+    if VERSION >= v"1.1"
+        # required ldiv!(triag, adjoint) that is not implemented in 1.0
+        register(:VECTOR_TO_VECTOR_FUNCS, [utriag_ldiv, ltriag_ldiv])
+        register(:MATRIX_TO_MATRIX_FUNCS, [utriag_ldiv, ltriag_ldiv])
+    end
+end
 
 end # module