Implement unconstraining transform for LKJCorr

[johncant]

I've ported this bijector from tensorflow and added to LKJCorr. This ensures that initial samples drawn from LKJCorr are positive definite, which fixes #7101 . Sampling now completes successfully with no divergences.

There are several parts I'm not comfortable with:

• https://github.com/johncant/pymc/blob/fix_lkjcorr_positive_definiteness/pymc/distributions/multivariate.py#L1583 - it seems like a bad idea to run a pytensor graph here. Is there any way to get the LKJCorr n parameter from op or rv without evaling any pytensors?
• https://github.com/johncant/pymc/blob/fix_lkjcorr_positive_definiteness/pymc/distributions/transforms.py#L176-L177 - not sure whether or not this is the right way to create a constant pytensor tensor here.

@fonnesbeck @twiecki @jessegrabowski @velochy - please could you take a look? I would like to make sure that this fix makes sense before adding tests and making the linters pass.

Notes:

• Tests not yet written, linters not yet ran
• The original tensorflow bijector is defined in the opposite sense to pymc transforms, i.e. forward in tensorflow_probability is backward in pymc
• The original tensorflow bijector produces cholesky factors, not actual correlation matrices, so in this implementation, we have to do a cholesky decomposition in the forward transform.
• In the tensorflow bijector, the triagonal elements of a matrix are filled in a clockwise spiral, as opposed to numpy which defines indices in a row-major order.

Description

Backward method

1. Start with identity matrix and fill lower triangular elements with unconstrained real numbers.
2. Normalize each row so the L-2 norm is 1
3. This is now a Cholesky factor that will always result in positive definite correlation matrices

Forward method

1. Reconstruct the correlation matrix from its upper triangular elements
2. Perform cholesky decomposition to obtain L
3. The diagonal elements of L are multipliers we used to normalize the other elements.
4. Extract those diagonal elements and divide to undo the backward method

log_jac_det

This was quite complicated to implement, so I used the symbolic jacobian.

Related Issue

• Closes BUG: LKJCorr breaks when used as covariance with MvNormal #7101

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

📚 Documentation preview 📚: https://pymc--7380.org.readthedocs.build/en/7380/

[welcome]

]
💖 Thanks for opening this pull request! 💖 The PyMC community really appreciates your time and effort to contribute to the project. Please make sure you have read our Contributing Guidelines and filled in our pull request template to the best of your ability.

[johncant]

Hi, It's unlikely I'm going to have any time to work on this for the next 6 months. The hardest part is coming up with a closed form solution for log_det_jac, which I don't think I'm very close to doing.

[twiecki]

Thanks for the update @johncant and for pushing this as far as you did.

[jessegrabowski]

This is updated to a fully working implementation of the transformer, with non-trivial tests.

It is currently blocked by #7910, because it uses pt.flip with negative indices, which are bugged in the current main branch. Tests pass locally on the current pytensor main.

There's some jank -- I'm not sure we need to pass n to the constructor. We can easily infer n from all the inputs. I guess the benefit is to be able to check that shapes haven't changed somewhere since the transformer was initialized? But these are so internal to PyMC that users shouldn't be able to make that happen anyway (and if one can, he's a super user so who cares).

If we got rid of the n property, it would also do away with the unit_diag arguments on the triangular_pack helpers.

On that note, those two triangular pack helpers could themselves be a separate transformation, since they define a Bijection with forward and backward methods. That might be conceptually better, but on the other hand I don't think these would be used anywhere else, so it also makes sense to keep them as part of this class.

[ricardoV94]

BTW once this is merged we should still explore the new transform by that Stan dev, it should have better sampling properties IIRC

[jessegrabowski]

Do you have a link to the actual implementation? I was rooting around in the stan/stan-math repo and couldn't find it

[johncant]

Wow - congratulations @jessegrabowski !

[jessegrabowski]

Wow - congratulations @jessegrabowski !

You had it 90% of the way there, we were just missing this weird spiral triangular construction that tfp was doing. I have no idea why they do it this way though, I just blindly copied xD

[jessegrabowski]

I switch the LKJCorr implementation to use scan to generate samples. This means:

1. Both eta and n can be symbolic
2. Larger matrices sample much faster.
3. But small matrices are somewhat slower

We could potentially try to check n and unroll the loop if it's sufficiently small. It also made me think an scan_to_unrolled_loop rewrite/tool might be interesting/useful in general.

Here is the benchmark I ran:

@pytest.mark.parametrize("n", [2, 10, 100])
def test_lkjcorr_sampling_benchmark(n, benchmark):
    d = pm.LKJCorr.dist(n=n, eta=2.0)
    sample_fn = pm.compile(inputs=[], outputs=[d], mode='FAST_RUN')
    sample_fn()

    benchmark(sample_fn)

Before (unrolled loop):

----------------------------------------------------------------------------------------------------- benchmark: 3 tests -----------------------------------------------------------------------------------------------------                                                                                                                                              
Name (time in us)                             Min                    Max                  Mean                StdDev                Median                 IQR            Outliers           OPS            Rounds  Iterations
------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------                                                                                                                                              
test_lkjcorr_sampling_benchmark[2]         3.1670 (1.0)          38.4999 (1.0)          3.6752 (1.0)          0.5252 (1.0)          3.7500 (1.0)        0.3759 (1.0)       435;361  272,093.1823 (1.0)       17765           1                                                                                                                                              
test_lkjcorr_sampling_benchmark[10]       59.1250 (18.67)       250.4170 (6.50)        68.1735 (18.55)        4.7520 (9.05)        69.0001 (18.40)      4.0832 (10.86)    2420;568   14,668.4544 (0.05)      12006           1
test_lkjcorr_sampling_benchmark[100]     831.8330 (262.66)   21,313.4581 (553.60)   1,588.8558 (432.32)   1,978.1464 (>1000.0)  1,050.2914 (280.08)   179.7080 (478.07)     63;117      629.3837 (0.00)        862           1                                                                                                                                              
------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------                                                                                                                                              

After (scan):

----------------------------------------------------------------------------------------------------- benchmark: 3 tests ----------------------------------------------------------------------------------------------------                                                                                                                                               
Name (time in us)                               Min                   Max                  Mean              StdDev                Median                 IQR            Outliers           OPS            Rounds  Iterations
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------                                                                                                                                               
test_lkjcorr_sampling_benchmark[2]           7.9999 (1.0)         77.4580 (1.0)          9.2234 (1.0)        2.3224 (1.0)          8.5421 (1.0)        0.9161 (1.0)       691;911  108,419.6575 (1.0)        8120           1                                                                                                                                               
test_lkjcorr_sampling_benchmark[10]        229.1250 (28.64)      581.1669 (7.50)       253.4111 (27.47)     21.7310 (9.36)       249.2500 (29.18)     19.0000 (20.74)     283;121    3,946.1577 (0.04)       3424           1
test_lkjcorr_sampling_benchmark[100]     3,908.5420 (488.57)   6,156.4590 (79.48)    4,609.0084 (499.71)   446.9420 (192.45)   4,542.3959 (531.77)   660.9370 (721.49)       70;3      216.9664 (0.00)        208           1                                                                                                                                               
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------                                                                                                                                               

[jessegrabowski]

We can support estimation of eta now. We should be able to add broadcasting on these parameters. but I'll leave that for a future PR.

[ricardoV94]

Those timings look good enough I wouldn't worry. It will hopefully fall out from our goal of speeding up Scan. Can you try with Numba mode out of curiosity?

[jessegrabowski]

Numba timings. With scan:

----------------------------------------------------------------------------------------------------- benchmark: 3 tests -----------------------------------------------------------------------------------------------------                                                                                                                                              
Name (time in us)                               Min                    Max                  Mean              StdDev                Median                 IQR            Outliers           OPS            Rounds  Iterations
------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------                                                                                                                                              
test_lkjcorr_sampling_benchmark[2]           8.5830 (1.0)          93.3340 (1.0)          9.6742 (1.0)        1.1092 (1.0)          9.8330 (1.0)        0.8750 (1.0)       215;186  103,367.4023 (1.0)       11911           1                                                                                                                                              
test_lkjcorr_sampling_benchmark[10]        199.0420 (23.19)    40,315.1250 (431.94)     272.4426 (28.16)    823.9673 (742.82)     225.4579 (22.93)     24.5001 (28.00)      23;326    3,670.4975 (0.04)       4113           1
test_lkjcorr_sampling_benchmark[100]     3,578.3750 (416.92)    5,509.3330 (59.03)    4,043.9366 (418.01)   256.9228 (231.62)   4,027.7915 (409.62)   316.2490 (361.44)       71;4      247.2838 (0.00)        238           1                                                                                                                                              
------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------                                                                                                                                              

Unrolled loop (old implementation):

-------------------------------------------------------------------------------------------------- benchmark: 3 tests --------------------------------------------------------------------------------------------------                                                                                                                                                    
Name (time in us)                             Min                   Max                Mean              StdDev              Median                 IQR             Outliers  OPS (Kops/s)            Rounds  Iterations
------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------                                                                                                                                                    
test_lkjcorr_sampling_benchmark[2]         4.7911 (1.0)         49.6670 (1.0)        5.5282 (1.0)        0.7295 (1.0)        5.5420 (1.0)        0.1250 (1.0)       174;2127      180.8892 (1.0)       10846           1                                                                                                                                                    
test_lkjcorr_sampling_benchmark[10]        7.5829 (1.58)        54.9171 (1.11)       8.7685 (1.59)       0.8032 (1.10)       8.8330 (1.59)       0.2501 (2.00)     4005;7133      114.0448 (0.63)      34385           1
test_lkjcorr_sampling_benchmark[100]     167.4580 (34.95)    2,400.6250 (48.33)    317.5335 (57.44)    235.1336 (322.34)   250.0410 (45.12)    105.5831 (844.46)     162;312        3.1493 (0.02)       2211           1                                                                                                                                                    
------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------                                                                                                                                                    

Pretty bad!

[jessegrabowski]

JAX also doesn't work with the new scan implementation, I guess because n is symbolic so pt.arange(2, n)(the length of the scan) is unknown at compile time

[ricardoV94]

length shouldn't matter if only last state is used? can you write index as a recurring out?

I don't recall where arange shows up

[ricardoV94]

timings are pretty nice, not worried about the small case

[jessegrabowski]

timings are pretty nice, not worried about the small case

No, it's bad. I updated the original timings with a compiled function instead of pm.draw and the scan is significantly worse across the board. I can't reproduce the original timings, so I might have been picking up first-run compile times or something.

[ricardoV94]

In the compiled version did you set updates so the rng is not copied?

[jessegrabowski]

I used pm.compile, doesn't that handle it for you?

[ricardoV94]

Ah so you were seeing longer compile times that made it look like scan was doing better? Unrolled 100 things or more is gonna be pretty ugly though

[ricardoV94]

Can you show the final numba graph?