Merge pull request #217 from JuliaStats/dw/three_arg_mul_ldiv

Rewrite `whiten!` and `unwhiten!` to work around ForwardDiff bug

Author: andreasnoack

Project.toml

@@ -1,6 +1,6 @@
 name = "PDMats"
 uuid = "90014a1f-27ba-587c-ab20-58faa44d9150"
-version = "0.11.33"
+version = "0.11.34"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -9,16 +9,20 @@ SuiteSparse = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"
 
 [compat]
 BandedMatrices = "0.15, 1"
+FiniteDifferences = "0.12"
+ForwardDiff = "0.10, 1"
 Random = "<0.0.1, 1"
 StaticArrays = "1"
 Test = "<0.0.1, 1"
 julia = "1.10"
 
 [extras]
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["BandedMatrices", "StaticArrays", "Random", "Test"]
+test = ["BandedMatrices", "FiniteDifferences", "ForwardDiff", "StaticArrays", "Random", "Test"]

src/chol.jl

@@ -61,8 +61,7 @@ function quad!(r::AbstractArray, A::Cholesky, X::AbstractMatrix)
     aU = chol_upper(A)
     z = similar(r, size(A, 1)) # buffer to save allocations
     @inbounds for i in axes(X, 2)
-        copyto!(z, view(X, :, i))
-        lmul!(aU, z)
+        mul!(z, aU, view(X, :, i))
         r[i] = sum(abs2, z)
     end
     return r

src/pdmat.jl

@@ -125,14 +125,20 @@ LinearAlgebra.sqrt(A::PDMat) = PDMat(sqrt(Hermitian(A.mat)))
 function whiten!(r::AbstractVecOrMat, a::PDMat, x::AbstractVecOrMat)
     @check_argdims axes(r) == axes(x)
     @check_argdims a.dim == size(x, 1)
-    v = _rcopy!(r, x)
-    return ldiv!(chol_lower(cholesky(a)), v)
+    if r === x
+        return ldiv!(chol_lower(cholesky(a)), r)
+    else
+        return ldiv!(r, chol_lower(cholesky(a)), x)
+    end
 end
 function unwhiten!(r::AbstractVecOrMat, a::PDMat, x::AbstractVecOrMat)
     @check_argdims axes(r) == axes(x)
     @check_argdims a.dim == size(x, 1)
-    v = _rcopy!(r, x)
-    return lmul!(chol_lower(cholesky(a)), v)
+    if r === x
+        return lmul!(chol_lower(cholesky(a)), r)
+    else
+        return mul!(r, chol_lower(cholesky(a)), x)
+    end
 end
 
 function whiten(a::PDMat, x::AbstractVecOrMat)
@@ -162,8 +168,7 @@ function quad!(r::AbstractArray, a::PDMat, x::AbstractMatrix)
     aU = chol_upper(cholesky(a))
     z = similar(r, a.dim) # buffer to save allocations
     @inbounds for i in axes(x, 2)
-        copyto!(z, view(x, :, i))
-        lmul!(aU, z)
+        mul!(z, aU, view(x, :, i))
         r[i] = sum(abs2, z)
     end
     return r

src/utils.jl

@@ -7,9 +7,6 @@ macro check_argdims(cond)
     end
 end
 
-_rcopy!(r, x) = (r === x || copyto!(r, x); r)
-
-
 function _addscal!(r::Matrix, a::Matrix, b::Union{Matrix, SparseMatrixCSC}, c::Real)
     if c == one(c)
         for i in eachindex(a)

test/ad.jl

@@ -0,0 +1,27 @@
+using PDMats
+using FiniteDifferences
+using ForwardDiff
+
+using LinearAlgebra
+using Test
+
+# issue #217
+@testset "PDMat: (un)whiten" begin
+    a = vec(Matrix{Float64}(I, 4, 4))
+    fdm = central_fdm(5, 1)
+
+    for (f, f!) in ((whiten, whiten!), (unwhiten, unwhiten!))
+        apply_f = let f = f
+            a -> f(PDMat(Symmetric(reshape(a, 4, 4))), ones(4))
+        end
+        apply_f! = let f! = f!
+            a -> f!(Vector{promote_type(eltype(a),Float64)}(undef, 4), PDMat(Symmetric(reshape(a, 4, 4))), ones(4))
+        end
+
+        J = only(FiniteDifferences.jacobian(fdm, apply_f, a))
+        @test only(FiniteDifferences.jacobian(fdm, apply_f!, a)) ≈ J
+    
+        @test ForwardDiff.jacobian(apply_f, a) ≈ J
+        @test ForwardDiff.jacobian(apply_f!, a) ≈ J
+    end
+end

test/runtests.jl

@@ -1,5 +1,5 @@
 include("testutils.jl")
-tests = ["pdmtypes", "abstracttypes", "addition", "generics", "kron", "chol", "specialarrays", "sqrt"]
+tests = ["pdmtypes", "abstracttypes", "addition", "generics", "kron", "chol", "specialarrays", "sqrt", "ad"]
 println("Running tests ...")
 
 for t in tests