Improve diag_At_A and diag_At_B efficiency (#288)

* Improve diag_At_A and diag_At_B for vectors

* Use dot

* Use sum for diag_At_A

* Add allocated tests

* Fix function signature

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Update tests to work with julia 1.3

* Patch release

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: sharanry

Project.toml

@@ -1,7 +1,7 @@
 name = "AbstractGPs"
 uuid = "99985d1d-32ba-4be9-9821-2ec096f28918"
 authors = ["JuliaGaussianProcesses Team"]
-version = "0.5.7"
+version = "0.5.8"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/util/common_covmat_ops.jl

@@ -60,18 +60,20 @@ Xt_invA_Y(X::AbstractVecOrMat, A::Cholesky, Y::AbstractVecOrMat) = (A.U' \ X)' *
 
 At_A(A::AbstractVecOrMat) = A'A
 
-diag_At_A(A::AbstractVecOrMat) = vec(sum(abs2.(A); dims=1))
+diag_At_A(A::AbstractMatrix) = vec(sum(abs2, A; dims=1))
+diag_At_A(A::AbstractVector) = [sum(abs2, A)]
 
 tr_At_A(A::AbstractVecOrMat) = sum(abs2, A)
 
-function diag_At_B(A::AbstractVecOrMat, B::AbstractVecOrMat)
+function diag_At_B(A::AbstractMatrix, B::AbstractMatrix)
     size(A) == size(B) || throw(
         DimensionMismatch(
             "A ($(size(A))) and B ($(size(B))) do not have the same dimensions "
         ),
     )
     return vec(sum(A .* B; dims=1))
 end
+diag_At_B(x::AbstractVector, y::AbstractVector) = [dot(x, y)]
 
 diag_Xt_A_X(A::Cholesky, X::AbstractVecOrMat) = diag_At_A(A.U * X)
 

test/util/common_covmat_ops.jl

@@ -76,9 +76,11 @@
     @test Xt_invA_Y(X, A, Y) ≈ X' * (A \ Y)
 
     @test diag_At_A(x) ≈ [x'x]
+    @test (@allocated diag_At_A(x)) <= 96
     @test diag_At_A(X) ≈ diag(X'X)
 
     @test diag_At_B(x, z) ≈ [x'z]
+    @test (@allocated diag_At_B(x, z)) <= 96
     @test diag_At_B(X, Z) ≈ diag(X'Z)
 
     @test diag_Xt_A_X(A, x) ≈ [Xt_A_X(A, x)]