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

📚 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

Attention: 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 20 commits behind head on main.

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

❌ Your patch status has failed because the patch coverage (70.37%) is below the target coverage (100.00%). You can increase the patch coverage or adjust the target coverage.

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

Open

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

Fyrebright commented Jun 19, 2025 •

edited

Loading

Uh oh!

ricardoV94 left a comment •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Fyrebright commented Jun 20, 2025

Uh oh!

jessegrabowski left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jessegrabowski commented Jun 21, 2025

Uh oh!

jessegrabowski commented Jun 22, 2025 •

edited

Loading

Uh oh!

ricardoV94 left a comment

Uh oh!

ricardoV94 Jun 22, 2025

Uh oh!

Fyrebright commented Jun 23, 2025

Uh oh!

codecov bot commented Jun 23, 2025 •

edited

Loading

Uh oh!

jessegrabowski left a comment

Uh oh!

Uh oh!

Uh oh!

	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

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

Uh oh!

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!

Uh oh!

Fyrebright commented Jun 20, 2025

Uh oh!

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jessegrabowski commented Jun 21, 2025

Uh oh!

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

Uh oh!

Uh oh!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Jun 22, 2025

Choose a reason for hiding this comment

Uh oh!

Fyrebright commented Jun 23, 2025

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

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