Use lapack func instead of `scipy.linalg.cholesky` #1487

Fyrebright · 2025-06-19T21:43:10Z

Now skips 2D checks in perform
Updated the default arguments for check_finite to false to match documentation
Add benchmark test case

Description

Adds _cholesky method to slinalg.Cholesky to replace the scipy.linalg.cholesky wrapper. It is almost identical to the corresponding scipy function but it skips the 2d check and the batching wrapper.

Previous performance (with check_finite=False):

benchmark: 5.1.0 (defaults: timer=time.perf_counter disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000) --------------------------------------------------- benchmark: 1 tests --------------------------------------------------- Name (time in us) Min Max Mean StdDev Median IQR Outliers OPS (Kops/s) Rounds Iterations -------------------------------------------------------------------------------------------------------------------------- test_cholesky_performance 5.6610 43.2910 6.8433 2.4983 5.9610 0.1810 1312;1574 146.1280 9487 1 --------------------------------------------------------------------------------------------------------------------------

After changes:

benchmark: 5.1.0 (defaults: timer=time.perf_counter disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000) --------------------------------------------------- benchmark: 1 tests -------------------------------------------------- Name (time in us) Min Max Mean StdDev Median IQR Outliers OPS (Kops/s) Rounds Iterations ------------------------------------------------------------------------------------------------------------------------- test_cholesky_performance 5.1500 28.5640 5.5494 1.0101 5.4210 0.1100 262;628 180.2012 12753 1 -------------------------------------------------------------------------------------------------------------------------

Related Issue

Closes #
Related to Directly use scipy lapack functions instead of scipy wrappers #1468

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): performance

📚 Documentation preview 📚: https://pytensor--1487.org.readthedocs.build/en/1487/

* Now skips 2D checks in perform * Updated the default arguments for `check_finite` to false to match documentation * Add benchmark test case

ricardoV94

Nice start, left some comments

pytensor/tensor/slinalg.py

Fyrebright · 2025-06-20T19:37:11Z

Thank you for the feedback. I moved it all to perform and removed the asarray,_datacopied checks. I also moved the empty input case to the beginning and added test coverage for it.

jessegrabowski

Looking really great! I see a few more places we can improve the performance (avoiding copy + cleaning up the code) then I think it'll be ready to merge

pytensor/tensor/slinalg.py

jessegrabowski · 2025-06-21T20:26:02Z

Looks really nice! I triggered a CI run, waiting to see if tests pass then will approve, pending thumbs up from @ricardoV94

jessegrabowski · 2025-06-22T10:40:21Z

You have a test failure in tests/link/numba/test_slinalg.py::test_cholesky_raises_on_nan_input:

2025-06-21T21:14:33.0385589Z =================================== FAILURES =================================== 2025-06-21T21:14:33.0386251Z ______________________ test_cholesky_raises_on_nan_input _______________________ 2025-06-21T21:14:33.0386688Z 2025-06-21T21:14:33.0386858Z def test_cholesky_raises_on_nan_input(): 2025-06-21T21:14:33.0387381Z test_value = rng.random(size=(3, 3)).astype(floatX) 2025-06-21T21:14:33.0387886Z test_value[0, 0] = np.nan 2025-06-21T21:14:33.0388256Z 2025-06-21T21:14:33.0388568Z x = pt.tensor(dtype=floatX, shape=(3, 3)) 2025-06-21T21:14:33.0389299Z x = x.T.dot(x) 2025-06-21T21:14:33.0389718Z g = pt.linalg.cholesky(x) 2025-06-21T21:14:33.0390184Z f = pytensor.function([x], g, mode="NUMBA") 2025-06-21T21:14:33.0390622Z 2025-06-21T21:14:33.0391099Z > with pytest.raises(np.linalg.LinAlgError, match=r"Non-numeric values"): 2025-06-21T21:14:33.0391737Z ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 2025-06-21T21:14:33.0392394Z E Failed: DID NOT RAISE <class 'numpy.linalg.LinAlgError'> 2025-06-21T21:14:33.0392804Z 2025-06-21T21:14:33.0392979Z tests/link/numba/test_slinalg.py:471: Failed

Could be related to the change in the default check_finite flag? You might need to adjust this test to correctly set the flag if so.

ricardoV94

Looks good. small comment

ricardoV94 · 2025-06-22T10:42:43Z

pytensor/tensor/slinalg.py

+
+ if self.check_finite and not np.isfinite(x).all():
+ if self.on_error == "nan":
+ out[0] = np.full(x.shape, np.nan, dtype=node.outputs[0].type.dtype)


Suggested change

out[0] = np.full(x.shape, np.nan, dtype=node.outputs[0].type.dtype)

out[0] = np.full(x.shape, np.nan, dtype=potrf.dtype)

Fyrebright · 2025-06-23T00:32:13Z

Could be related to the change in the default check_finite flag? You might need to adjust this test to correctly set the flag if so.

It is fixed now. I had been skipping the link/* tests until now. They are all passing on my machine except for:

FAILED tests/scan/test_basic.py::test_cython_performance - ModuleNotFoundError: No module named 'pytensor.scan.scan_perform' FAILED tests/scan/test_basic.py::test_output_storage_reuse[cvm] - ModuleNotFoundError: No module named 'pytensor.scan.scan_perform'

codecov · 2025-06-23T07:30:10Z

Codecov Report

❌ Patch coverage is 70.37037% with 8 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@d3bbc20). Learn more about missing BASE report.
⚠️ Report is 138 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/tensor/slinalg.py	70.37%	5 Missing and 3 partials ⚠️

Additional details and impacted files

@@ Coverage Diff @@ ## main #1487 +/- ## ======================================= Coverage ? 81.98% ======================================= Files ? 231 Lines ? 52192 Branches ? 9185 ======================================= Hits ? 42790 Misses ? 7094 Partials ? 2308

Files with missing lines	Coverage Δ
pytensor/tensor/slinalg.py	`92.19% <70.37%> (ø)`

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

jessegrabowski

Looks great, thanks for this!

Use lapack func instead of scipy.linalg.cholesky

98fc6a7

* Now skips 2D checks in perform * Updated the default arguments for `check_finite` to false to match documentation * Add benchmark test case

Fyrebright marked this pull request as ready for review June 19, 2025 22:09

ricardoV94 reviewed Jun 20, 2025

View reviewed changes

pytensor/tensor/slinalg.py Outdated Show resolved Hide resolved

pytensor/tensor/slinalg.py Outdated Show resolved Hide resolved

pytensor/tensor/slinalg.py Outdated Show resolved Hide resolved

ricardoV94 added maintenance performance labels Jun 20, 2025

Refactor out _cholesky helper, add empty test

1af5b61

jessegrabowski requested changes Jun 21, 2025

View reviewed changes

pytensor/tensor/slinalg.py Outdated Show resolved Hide resolved

pytensor/tensor/slinalg.py Outdated Show resolved Hide resolved

pytensor/tensor/slinalg.py Outdated Show resolved Hide resolved

pytensor/tensor/slinalg.py Outdated Show resolved Hide resolved

jessegrabowski mentioned this pull request Jun 21, 2025

Handle lower flag more carefully in Solve to CholeskySolve rewrite #1492

Closed

11 tasks

Remove array and potrf copies

2c600d9

Fyrebright requested a review from jessegrabowski June 21, 2025 20:14

ricardoV94 reviewed Jun 22, 2025

View reviewed changes

Update test_cholesky_raises_on_nan_input

f0435b1

ricardoV94 approved these changes Jun 23, 2025

View reviewed changes

jessegrabowski approved these changes Jun 23, 2025

View reviewed changes

jessegrabowski merged commit 236e50d into pymc-devs:main Jun 23, 2025
73 of 74 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Use lapack func instead of `scipy.linalg.cholesky` #1487

Use lapack func instead of `scipy.linalg.cholesky` #1487

Uh oh!

Fyrebright commented Jun 19, 2025 •

edited

Loading

ricardoV94 left a comment •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Fyrebright commented Jun 20, 2025

jessegrabowski left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jessegrabowski commented Jun 21, 2025

jessegrabowski commented Jun 22, 2025 •

edited

Loading

ricardoV94 left a comment

ricardoV94 Jun 22, 2025

Fyrebright commented Jun 23, 2025

codecov bot commented Jun 23, 2025 •

edited

Loading

jessegrabowski left a comment

Uh oh!

Labels

3 participants

	out[0] = np.full(x.shape, np.nan, dtype=node.outputs[0].type.dtype)
	out[0] = np.full(x.shape, np.nan, dtype=potrf.dtype)

Use lapack func instead of scipy.linalg.cholesky #1487

Use lapack func instead of scipy.linalg.cholesky #1487

Uh oh!

Conversation

Fyrebright commented Jun 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue

Checklist

Type of change

ricardoV94 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Fyrebright commented Jun 20, 2025

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jessegrabowski commented Jun 21, 2025

jessegrabowski commented Jun 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

ricardoV94 left a comment

Choose a reason for hiding this comment

ricardoV94 Jun 22, 2025

Choose a reason for hiding this comment

Fyrebright commented Jun 23, 2025

codecov bot commented Jun 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

Labels

3 participants

Use lapack func instead of `scipy.linalg.cholesky` #1487

Use lapack func instead of `scipy.linalg.cholesky` #1487

Fyrebright commented Jun 19, 2025 •

edited

Loading

ricardoV94 left a comment •

edited

Loading

jessegrabowski commented Jun 22, 2025 •

edited

Loading

codecov bot commented Jun 23, 2025 •

edited

Loading