Loosen AbstractArray restrictions where these are not necessary (#63)

* remove `AbstractArray` annotations for rrules

* remove `AbstractArray` annotations for implementations

Author: lkdvos

ext/MatrixAlgebraKitChainRulesCoreExt.jl

@@ -11,7 +11,7 @@ using LinearAlgebra
 
 MatrixAlgebraKit.iszerotangent(::AbstractZero) = true
 
-function ChainRulesCore.rrule(::typeof(copy_input), f, A::AbstractMatrix)
+function ChainRulesCore.rrule(::typeof(copy_input), f, A)
     project = ProjectTo(A)
     copy_input_pullback(ΔA) = (NoTangent(), NoTangent(), project(unthunk(ΔA)))
     return copy_input(f, A), copy_input_pullback
@@ -20,7 +20,7 @@ end
 for qr_f in (:qr_compact, :qr_full)
     qr_f! = Symbol(qr_f, '!')
     @eval begin
-        function ChainRulesCore.rrule(::typeof($qr_f!), A::AbstractMatrix, QR, alg)
+        function ChainRulesCore.rrule(::typeof($qr_f!), A, QR, alg)
             Ac = copy_input($qr_f, A)
             QR = $(qr_f!)(Ac, QR, alg)
             function qr_pullback(ΔQR)
@@ -58,7 +58,7 @@ end
 for lq_f in (:lq_compact, :lq_full)
     lq_f! = Symbol(lq_f, '!')
     @eval begin
-        function ChainRulesCore.rrule(::typeof($lq_f!), A::AbstractMatrix, LQ, alg)
+        function ChainRulesCore.rrule(::typeof($lq_f!), A, LQ, alg)
             Ac = copy_input($lq_f, A)
             LQ = $(lq_f!)(Ac, LQ, alg)
             function lq_pullback(ΔLQ)
@@ -102,7 +102,7 @@ for eig in (:eig, :eigh)
     eig_t_pb = Symbol(eig, "_trunc_pullback")
     _make_eig_t_pb = Symbol("_make_", eig_t_pb)
     @eval begin
-        function ChainRulesCore.rrule(::typeof($eig_f!), A::AbstractMatrix, DV, alg)
+        function ChainRulesCore.rrule(::typeof($eig_f!), A, DV, alg)
             Ac = copy_input($eig_f, A)
             DV = $(eig_f!)(Ac, DV, alg)
             function $eig_pb(ΔDV)
@@ -115,10 +115,7 @@ for eig in (:eig, :eigh)
             end
             return DV, $eig_pb
         end
-        function ChainRulesCore.rrule(
-                ::typeof($eig_t!), A::AbstractMatrix, DV,
-                alg::TruncatedAlgorithm
-            )
+        function ChainRulesCore.rrule(::typeof($eig_t!), A, DV, alg::TruncatedAlgorithm)
             Ac = copy_input($eig_f, A)
             DV = $(eig_f!)(Ac, DV, alg.alg)
             DV′, ind = MatrixAlgebraKit.truncate($eig_t!, DV, alg.trunc)
@@ -141,7 +138,7 @@ end
 for svd_f in (:svd_compact, :svd_full)
     svd_f! = Symbol(svd_f, "!")
     @eval begin
-        function ChainRulesCore.rrule(::typeof($svd_f!), A::AbstractMatrix, USVᴴ, alg)
+        function ChainRulesCore.rrule(::typeof($svd_f!), A, USVᴴ, alg)
             Ac = copy_input($svd_f, A)
             USVᴴ = $(svd_f!)(Ac, USVᴴ, alg)
             function svd_pullback(ΔUSVᴴ)
@@ -157,10 +154,7 @@ for svd_f in (:svd_compact, :svd_full)
     end
 end
 
-function ChainRulesCore.rrule(
-        ::typeof(svd_trunc!), A::AbstractMatrix, USVᴴ,
-        alg::TruncatedAlgorithm
-    )
+function ChainRulesCore.rrule(::typeof(svd_trunc!), A, USVᴴ, alg::TruncatedAlgorithm)
     Ac = copy_input(svd_compact, A)
     USVᴴ = svd_compact!(Ac, USVᴴ, alg.alg)
     USVᴴ′, ind = MatrixAlgebraKit.truncate(svd_trunc!, USVᴴ, alg.trunc)
@@ -178,7 +172,7 @@ function _make_svd_trunc_pullback(A, USVᴴ, ind)
     return svd_trunc_pullback
 end
 
-function ChainRulesCore.rrule(::typeof(left_polar!), A::AbstractMatrix, WP, alg)
+function ChainRulesCore.rrule(::typeof(left_polar!), A, WP, alg)
     Ac = copy_input(left_polar, A)
     WP = left_polar!(Ac, WP, alg)
     function left_polar_pullback(ΔWP)
@@ -192,7 +186,7 @@ function ChainRulesCore.rrule(::typeof(left_polar!), A::AbstractMatrix, WP, alg)
     return WP, left_polar_pullback
 end
 
-function ChainRulesCore.rrule(::typeof(right_polar!), A::AbstractMatrix, PWᴴ, alg)
+function ChainRulesCore.rrule(::typeof(right_polar!), A, PWᴴ, alg)
     Ac = copy_input(left_polar, A)
     PWᴴ = right_polar!(Ac, PWᴴ, alg)
     function right_polar_pullback(ΔPWᴴ)

src/implementations/eig.jl

@@ -3,9 +3,7 @@
 function copy_input(::typeof(eig_full), A::AbstractMatrix)
     return copy!(similar(A, float(eltype(A))), A)
 end
-function copy_input(::typeof(eig_vals), A)
-    return copy_input(eig_full, A)
-end
+copy_input(::typeof(eig_vals), A) = copy_input(eig_full, A)
 copy_input(::typeof(eig_trunc), A) = copy_input(eig_full, A)
 
 copy_input(::typeof(eig_full), A::Diagonal) = copy(A)
@@ -67,7 +65,7 @@ function initialize_output(::typeof(eig_vals!), A::AbstractMatrix, ::AbstractAlg
     D = similar(A, Tc, n)
     return D
 end
-function initialize_output(::typeof(eig_trunc!), A::AbstractMatrix, alg::TruncatedAlgorithm)
+function initialize_output(::typeof(eig_trunc!), A, alg::TruncatedAlgorithm)
     return initialize_output(eig_full!, A, alg.alg)
 end
 
@@ -108,7 +106,7 @@ function eig_vals!(A::AbstractMatrix, D, alg::LAPACK_EigAlgorithm)
     return D
 end
 
-function eig_trunc!(A::AbstractMatrix, DV, alg::TruncatedAlgorithm)
+function eig_trunc!(A, DV, alg::TruncatedAlgorithm)
     D, V = eig_full!(A, DV, alg.alg)
     return first(truncate(eig_trunc!, (D, V), alg.trunc))
 end
@@ -131,7 +129,7 @@ end
 
 # GPU logic
 # ---------
-_gpu_geev!(A::AbstractMatrix, D, V) = throw(MethodError(_gpu_geev!, (A, D, V)))
+_gpu_geev!(A, D, V) = throw(MethodError(_gpu_geev!, (A, D, V)))
 
 function eig_full!(A::AbstractMatrix, DV, alg::GPU_EigAlgorithm)
     check_input(eig_full!, A, DV, alg)

src/implementations/eigh.jl

@@ -3,9 +3,7 @@
 function copy_input(::typeof(eigh_full), A::AbstractMatrix)
     return copy!(similar(A, float(eltype(A))), A)
 end
-function copy_input(::typeof(eigh_vals), A::AbstractMatrix)
-    return copy_input(eigh_full, A)
-end
+copy_input(::typeof(eigh_vals), A) = copy_input(eigh_full, A)
 copy_input(::typeof(eigh_trunc), A) = copy_input(eigh_full, A)
 
 copy_input(::typeof(eigh_full), A::Diagonal) = copy(A)
@@ -65,7 +63,7 @@ function initialize_output(::typeof(eigh_vals!), A::AbstractMatrix, ::AbstractAl
     D = similar(A, real(eltype(A)), n)
     return D
 end
-function initialize_output(::typeof(eigh_trunc!), A::AbstractMatrix, alg::TruncatedAlgorithm)
+function initialize_output(::typeof(eigh_trunc!), A, alg::TruncatedAlgorithm)
     return initialize_output(eigh_full!, A, alg.alg)
 end
 
@@ -111,7 +109,7 @@ function eigh_vals!(A::AbstractMatrix, D, alg::LAPACK_EighAlgorithm)
     return D
 end
 
-function eigh_trunc!(A::AbstractMatrix, DV, alg::TruncatedAlgorithm)
+function eigh_trunc!(A, DV, alg::TruncatedAlgorithm)
     D, V = eigh_full!(A, DV, alg.alg)
     return first(truncate(eigh_trunc!, (D, V), alg.trunc))
 end

src/implementations/gen_eig.jl

@@ -3,9 +3,7 @@
 function copy_input(::typeof(gen_eig_full), A::AbstractMatrix, B::AbstractMatrix)
     return copy!(similar(A, float(eltype(A))), A), copy!(similar(B, float(eltype(B))), B)
 end
-function copy_input(::typeof(gen_eig_vals), A::AbstractMatrix, B::AbstractMatrix)
-    return copy_input(gen_eig_full, A, B)
-end
+copy_input(::typeof(gen_eig_vals), A, B) = copy_input(gen_eig_full, A, B)
 
 function check_input(::typeof(gen_eig_full!), A::AbstractMatrix, B::AbstractMatrix, WV, ::AbstractAlgorithm)
     ma, na = size(A)

src/implementations/lq.jl

@@ -1,11 +1,10 @@
 # Inputs
 # ------
-for f in (:lq_full, :lq_compact, :lq_null)
-    @eval function copy_input(::typeof($f), A)
-        return copy!(similar(A, float(eltype(A))), A)
-    end
-    @eval copy_input(::typeof($f), A::Diagonal) = copy(A)
-end
+copy_input(::typeof(lq_full), A::AbstractMatrix) = copy!(similar(A, float(eltype(A))), A)
+copy_input(::typeof(lq_compact), A) = copy_input(lq_full, A)
+copy_input(::typeof(lq_null), A) = copy_input(lq_full, A)
+
+copy_input(::typeof(lq_full), A::Diagonal) = copy(A)
 
 function check_input(::typeof(lq_full!), A::AbstractMatrix, LQ, ::AbstractAlgorithm)
     m, n = size(A)

src/implementations/qr.jl

@@ -1,11 +1,10 @@
 # Inputs
 # ------
-for f in (:qr_full, :qr_compact, :qr_null)
-    @eval function copy_input(::typeof($f), A)
-        return copy!(similar(A, float(eltype(A))), A)
-    end
-    @eval copy_input(::typeof($f), A::Diagonal) = copy(A)
-end
+copy_input(::typeof(qr_full), A::AbstractMatrix) = copy!(similar(A, float(eltype(A))), A)
+copy_input(::typeof(qr_compact), A) = copy_input(qr_full, A)
+copy_input(::typeof(qr_null), A) = copy_input(qr_full, A)
+
+copy_input(::typeof(qr_full), A::Diagonal) = copy(A)
 
 function check_input(::typeof(qr_full!), A::AbstractMatrix, QR, ::AbstractAlgorithm)
     m, n = size(A)

src/implementations/schur.jl

@@ -1,7 +1,7 @@
 # Inputs
 # ------
-copy_input(::typeof(schur_full), A::AbstractMatrix) = copy_input(eig_full, A)
-copy_input(::typeof(schur_vals), A::AbstractMatrix) = copy_input(eig_vals, A)
+copy_input(::typeof(schur_full), A) = copy_input(eig_full, A)
+copy_input(::typeof(schur_vals), A) = copy_input(eig_vals, A)
 
 # check input
 function check_input(::typeof(schur_full!), A::AbstractMatrix, TZv, ::AbstractAlgorithm)

src/implementations/svd.jl

@@ -1,8 +1,6 @@
 # Inputs
 # ------
-function copy_input(::typeof(svd_full), A)
-    return copy!(similar(A, float(eltype(A))), A)
-end
+copy_input(::typeof(svd_full), A::AbstractMatrix) = copy!(similar(A, float(eltype(A))), A)
 copy_input(::typeof(svd_compact), A) = copy_input(svd_full, A)
 copy_input(::typeof(svd_vals), A) = copy_input(svd_full, A)
 copy_input(::typeof(svd_trunc), A) = copy_input(svd_compact, A)
@@ -238,7 +236,7 @@ function svd_vals!(A::AbstractMatrix, S, alg::LAPACK_SVDAlgorithm)
     return S
 end
 
-function svd_trunc!(A::AbstractMatrix, USVᴴ, alg::TruncatedAlgorithm)
+function svd_trunc!(A, USVᴴ, alg::TruncatedAlgorithm)
     USVᴴ′ = svd_compact!(A, USVᴴ, alg.alg)
     return first(truncate(svd_trunc!, USVᴴ′, alg.trunc))
 end
@@ -270,7 +268,7 @@ function svd_full!(A::AbstractMatrix, USVᴴ, alg::DiagonalAlgorithm)
 
     return U, S, Vᴴ
 end
-function svd_compact!(A::AbstractMatrix, USVᴴ, alg::DiagonalAlgorithm)
+function svd_compact!(A, USVᴴ, alg::DiagonalAlgorithm)
     return svd_full!(A, USVᴴ, alg)
 end
 function svd_vals!(A::AbstractMatrix, S, alg::DiagonalAlgorithm)

src/implementations/truncation.jl

@@ -29,9 +29,7 @@ end
 # findtruncated
 # -------------
 # Generic fallback
-function findtruncated_svd(values::AbstractVector, strategy::TruncationStrategy)
-    return findtruncated(values, strategy)
-end
+findtruncated_svd(values, strategy::TruncationStrategy) = findtruncated(values, strategy)
 
 # specific implementations for finding truncated values
 findtruncated(values::AbstractVector, ::NoTruncation) = Colon()