AD fix for CorrBijector

[torfjelde]

This is a sibling-PR of #280, making use of the functionality introduced there to ensure that CorrBijector and it's siblings are also working as intended.

[yebai]

There is a reproducible Cholesky factorization failure error. I couldn't find any obvious sources causing it.

[sunxd3]

Looks like in v1.6 https://github.com/JuliaLang/julia/blob/3b76b25b648cc1fa6187b118c685819a4111a5d2/stdlib/LinearAlgebra/src/cholesky.jl#L377C1-L377C1 and https://github.com/JuliaLang/julia/blob/3b76b25b648cc1fa6187b118c685819a4111a5d2/stdlib/LinearAlgebra/src/cholesky.jl#L401 have too stringent type signiture.

In latest version cholesky takes AbstractMatrix.

[sunxd3]

Also when I try to run https://github.com/TuringLang/Bijectors.jl/blob/06b936d334bc2e58759ca5bbb830e225e8703b5a/test/transform.jl#L216-L230

encounter

julia> dist = LKJCholesky(3, 1)
LKJCholesky{Float64}(
d: 3
η: 1.0
uplo: L
)

julia> x = rand(dist)
Cholesky{Float64, Matrix{Float64}}
L factor:
3×3 LowerTriangular{Float64, Matrix{Float64}}:
  1.0        ⋅         ⋅ 
  0.198156  0.980171   ⋅ 
 -0.797754  0.257329  0.545316

julia> J = ForwardDiff.jacobian(x -> link(dist, x), x.U)
ERROR: PosDefException: matrix is not positive definite; Cholesky factorization failed.
Stacktrace:
  [1] checkpositivedefinite
    @ /path/to/julia/stdlib/v1.9/LinearAlgebra/src/factorization.jl:18 [inlined]
  [2] #cholesky!#152
    @ /path/to/julia/stdlib/v1.9/LinearAlgebra/src/cholesky.jl:268 [inlined]
  [3] cholesky! (repeats 2 times)
    @ /path/to/julia/stdlib/v1.9/LinearAlgebra/src/cholesky.jl:266 [inlined]
  [4] #cholesky#162
    @ /path/to/julia/stdlib/v1.9/LinearAlgebra/src/cholesky.jl:400 [inlined]
  [5] cholesky (repeats 2 times)
    @ /path/to/julia/stdlib/v1.9/LinearAlgebra/src/cholesky.jl:400 [inlined]
  [6] cholesky_lower(X::UpperTriangular{ForwardDiff.Dual{ForwardDiff.Tag{var"#9#10", Float64}, Float64, 9}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{var"#9#10", Float64}, Float64, 9}}})
    @ Bijectors /path/to/Bijectors/src/utils.jl:31
  [7] transform(b::Bijectors.VecCholeskyBijector, X::UpperTriangular{ForwardDiff.Dual{ForwardDiff.Tag{var"#9#10", Float64}, Float64, 9}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{var"#9#10", Float64}, Float64, 9}}})
    @ Bijectors /path/to/Bijectors/src/bijectors/corr.jl:220
  [8] Transform
    @ /path/to/Bijectors/src/interface.jl:80 [inlined]
  [9] link(d::LKJCholesky{Float64}, x::UpperTriangular{ForwardDiff.Dual{ForwardDiff.Tag{var"#9#10", Float64}, Float64, 9}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{var"#9#10", Float64}, Float64, 9}}})
    @ Bijectors /path/to/Bijectors/src/Bijectors.jl:128
 [10] (::var"#9#10")(x::UpperTriangular{ForwardDiff.Dual{ForwardDiff.Tag{var"#9#10", Float64}, Float64, 9}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{var"#9#10", Float64}, Float64, 9}}})
    @ Main ./REPL[3]:1
 [11] vector_mode_dual_eval!
    @ /path/to/ForwardDiff/src/apiutils.jl:24 [inlined]
 [12] vector_mode_jacobian(f::var"#9#10", x::UpperTriangular{Float64, Matrix{Float64}}, cfg::ForwardDiff.JacobianConfig{ForwardDiff.Tag{var"#9#10", Float64}, Float64, 9, UpperTriangular{ForwardDiff.Dual{ForwardDiff.Tag{var"#9#10", Float64}, Float64, 9}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{var"#9#10", Float64}, Float64, 9}}}})
    @ ForwardDiff /path/to/ForwardDiff/src/jacobian.jl:125
 [13] jacobian(f::Function, x::UpperTriangular{Float64, Matrix{Float64}}, cfg::ForwardDiff.JacobianConfig{ForwardDiff.Tag{var"#9#10", Float64}, Float64, 9, UpperTriangular{ForwardDiff.Dual{ForwardDiff.Tag{var"#9#10", Float64}, Float64, 9}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{var"#9#10", Float64}, Float64, 9}}}}, ::Val{true})
    @ ForwardDiff /path/to/ForwardDiff/src/jacobian.jl:21
 [14] jacobian(f::Function, x::UpperTriangular{Float64, Matrix{Float64}}, cfg::ForwardDiff.JacobianConfig{ForwardDiff.Tag{var"#9#10", Float64}, Float64, 9, UpperTriangular{ForwardDiff.Dual{ForwardDiff.Tag{var"#9#10", Float64}, Float64, 9}, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{var"#9#10", Float64}, Float64, 9}}}})
    @ ForwardDiff /path/to/ForwardDiff/src/jacobian.jl:19
 [15] jacobian(f::Function, x::UpperTriangular{Float64, Matrix{Float64}})
    @ ForwardDiff /path/to/ForwardDiff/src/jacobian.jl:19
 [16] top-level scope
    @ REPL[3]:1

[sunxd3]

Maybe related, LKJCholesky default uplo='L'; Hermitian default uplo='U'

[torfjelde]

In latest version cholesky takes AbstractMatrix.

But we only call cholesky with Hermitian, which is defined, no?

Maybe related, LKJCholesky default uplo='L'; Hermitian default uplo='U'

The issue is two-fold. First, we're giving the transformation in ForwardDiff.jacobian(...) a upper-triangular but it expects a lower-triangular (it ends up calling cholesky_lower). But even if we fix this, i.e. pass x.U instead, lit just returns 0s:

julia> using Bijectors

julia> dist = LKJCholesky(3, 1)
LKJCholesky{Float64}(
d: 3
η: 1.0
uplo: L
)

julia> x = rand(dist)
LinearAlgebra.Cholesky{Float64, Matrix{Float64}}
L factor:
3×3 LinearAlgebra.LowerTriangular{Float64, Matrix{Float64}}:
  1.0         ⋅         ⋅ 
 -0.107859   0.994166   ⋅ 
 -0.38706   -0.246285  0.888554

julia> link(dist, x)
3-element Vector{Float64}:
 -0.10828016710654258
 -0.40833723913774234
 -0.2737439074161078

julia> link(dist, x.L)
3-element Vector{Float64}:
 0.0
 0.0
 0.0

This was caused by this line: https://github.com/TuringLang/Bijectors.jl/blob/5bffcfb5a3c63b7515b3275acfdf742149dcbb04/src/utils.jl#L31

Here we ended up calling Hermitian(X, :U) on a X which is lower-triangular; I've now made this call Hermitian(X, :L). Now we have transpose(cholesky_lower(x.L)) == cholesky_upper(x.U), as expected.

[torfjelde]

Alrighty; tests are at least passing locally for me now. Let's see if Julia 1.6 also works.

[torfjelde]

Looks like we're ready to go:)

[yebai]

Thanks @sunxd3 and @torfjelde!