AD fix for CorrBijector #281

torfjelde · 2023-08-07T16:37:13Z

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.

torfjelde · 2023-08-07T16:37:52Z

ext/BijectorsZygoteExt.jl

+# TODO: Remove these as soon as https://github.com/FluxML/Zygote.jl/pull/1444 is merged.
+@adjoint LinearAlgebra.parent(x::LinearAlgebra.UpperTriangular) = parent(x), ȳ -> (LinearAlgebra.UpperTriangular(ȳ),)
+@adjoint LinearAlgebra.parent(x::LinearAlgebra.LowerTriangular) = parent(x), ȳ -> (LinearAlgebra.LowerTriangular(ȳ),)


I've added this here to check if tests are succeeding; will remove as soon as the mentioned PR goes through.

torfjelde · 2023-08-07T16:39:09Z

src/bijectors/corr.jl

@@ -321,7 +325,7 @@ function _link_chol_lkj(W::UpperTriangular)
 return z
 end

-_link_chol_lkj(W::LowerTriangular) = _link_chol_lkj(transpose(W))
+_link_chol_lkj_from_lower(W::AbstractMatrix) = _link_chol_lkj_from_upper(transpose_eager(W))


This one is a bit weird because I just saw that it has it's own chainrule defined while the evaluation itself does not have a specialized implementation..

ext/BijectorsReverseDiffExt.jl

yebai · 2023-08-10T14:51:56Z

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

* Update chainrules.jl * Update corr.jl * Revert changes to transform.

sunxd3 · 2023-08-11T13:20:31Z

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 · 2023-08-11T13:57:39Z

Also when I try to run

Bijectors.jl/test/transform.jl

Lines 216 to 230 in 06b936d

 @testset "LKJCholesky" begin 

 dist = LKJCholesky(3, 1) 

 single_sample_tests(dist) 

 x = rand(dist) 

 upperinds = [ 

 LinearIndices(size(x))[I] for I in CartesianIndices(size(x)) if I[2] > I[1] 

 ] 

 J = ForwardDiff.jacobian(x -> link(dist, x), x.U) 

 J = J[:, upperinds] 

 logpdf_turing = logpdf_with_trans(dist, x, true) 

 @test logpdf(dist, x) - _logabsdet(J) ≈ logpdf_turing 

 end

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 · 2023-08-11T14:02:00Z

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

version in test

…orfjelde/corr-ad-update

torfjelde · 2023-08-12T08:45:58Z

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:

Bijectors.jl/src/utils.jl

Line 31 in 5bffcfb

 cholesky_lower(X::AbstractMatrix) = lower_triangular(parent(cholesky(Hermitian(X)).L)) 

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.

in LKJCholesky tests

torfjelde · 2023-08-12T09:08:33Z

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

torfjelde · 2023-08-12T10:03:45Z

Looks like we're ready to go:)

@Grad

* added cholesky_lower and cholesky_triangular * updated PD to use new cholesky_lower and cholesky_upper * simplified imports in BijectorsReverseDiffExtx * added ChainRules as a dep since we need the chain rules for cholesky, etc. * forgot to update Project.toml in previous commit * added explicit implementation of with_logabsdet_jacobian for PDBijector * Update src/utils.jl * added ProjectTo in rrules for cholesky_lower and cholesky_upper to be proper * added ProjectTo for cholesky_upper too * added transpose_eager as a alias for permutedims to allow definition of AD rules without type piracy * allow usage of ForwardDiff gradient as ground-truth * added AD tests for PDVecBijector * added AD tests for PDVecBijector to runtests and commented out all other tests for the sake of reproducing ReverseDiff bug * forgot to remove type-piracy def of ReverseDiff rule for permutedims * use ReverseDiff.@Grad instead of ReverseDiff.@grad_from_chainrules * only define cholesky_lower and cholesky_upper rules for ReverseDiff, remove rules ChainRules defs * formatting * parameterise gradient test for PD bijector properly instead of using ForwardDiff as per suggestion of @devmotion * reversed chagne to test_ad * reactivate tests * updated doocstrings * improved PDVecBijector AD tests a bit * AD fix for CorrBijector (#281) * removed redundant imports to BijectorsZygoteExt * use cholesky_upper and cholesky_lower instead of cholesky_factor, etc. * added tests for CorrVecBijector * name testset correctly * use cholesky_lower and cholesky_upper instead of cholesky_factor * removed now-redundant cholesky_factor * Fix obsolete function references in tests. (#282) * Update chainrules.jl * Update corr.jl * Revert changes to transform. * removed type-piracy that has been addressed upstream and bumped Zygote version in test * use :L for Hermitian in `cholesky_lower` * fixed ForwardDiff tests for LKJCholesky * fixed tests for matrix dists and added tests for both values of uplo in LKJCholesky tests * another attempt at fixing Julia 1.6 tests --------- Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com> --------- Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com>

yebai · 2023-08-12T11:54:17Z

Thanks @sunxd3 and @torfjelde!

torfjelde added 5 commits August 7, 2023 11:45

removed redundant imports to BijectorsZygoteExt

7fcae80

use cholesky_upper and cholesky_lower instead of cholesky_factor, etc.

2e24cb8

added tests for CorrVecBijector

181e7d5

name testset correctly

09ca77c

use cholesky_lower and cholesky_upper instead of cholesky_factor

7c35448

torfjelde commented Aug 7, 2023

View reviewed changes

removed now-redundant cholesky_factor

5bffcfb

torfjelde commented Aug 7, 2023

View reviewed changes

torfjelde requested a review from devmotion August 7, 2023 16:42

yebai reviewed Aug 10, 2023

View reviewed changes

ext/BijectorsReverseDiffExt.jl Show resolved Hide resolved

Fix obsolete function references in tests. (#282)

06b936d

* Update chainrules.jl * Update corr.jl * Revert changes to transform.

torfjelde added 4 commits August 12, 2023 08:58

removed type-piracy that has been addressed upstream and bumped Zygote

8e16077

version in test

Merge remote-tracking branch 'origin/torfjelde/corr-ad-update' into t…

cec71b4

…orfjelde/corr-ad-update

use :L for Hermitian in cholesky_lower

51693bb

fixed ForwardDiff tests for LKJCholesky

182d868

fixed tests for matrix dists and added tests for both values of uplo

ae2370f

in LKJCholesky tests

another attempt at fixing Julia 1.6 tests

1555605

yebai approved these changes Aug 12, 2023

View reviewed changes

torfjelde merged commit 52ee210 into torfjelde/pd-fix Aug 12, 2023
21 checks passed

delete-merged-branch bot deleted the torfjelde/corr-ad-update branch August 12, 2023 10:47

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

AD fix for CorrBijector #281

AD fix for CorrBijector #281

torfjelde commented Aug 7, 2023

torfjelde Aug 7, 2023

torfjelde Aug 7, 2023

yebai commented Aug 10, 2023

sunxd3 commented Aug 11, 2023

sunxd3 commented Aug 11, 2023

sunxd3 commented Aug 11, 2023

torfjelde commented Aug 12, 2023

torfjelde commented Aug 12, 2023

torfjelde commented Aug 12, 2023

yebai commented Aug 12, 2023

AD fix for CorrBijector #281

AD fix for CorrBijector #281

Conversation

torfjelde commented Aug 7, 2023

torfjelde Aug 7, 2023

Choose a reason for hiding this comment

torfjelde Aug 7, 2023

Choose a reason for hiding this comment

yebai commented Aug 10, 2023

sunxd3 commented Aug 11, 2023

sunxd3 commented Aug 11, 2023

sunxd3 commented Aug 11, 2023

torfjelde commented Aug 12, 2023

torfjelde commented Aug 12, 2023

torfjelde commented Aug 12, 2023

yebai commented Aug 12, 2023