Use LAPACK functions for `cho_solve`, `lu_factor`, `solve_triangular` #1605

Fyrebright · 2025-09-02T15:53:16Z

Description

Directly use LAPACK functions in the perform method of the following classes, removing some checks from scipy.linalg (e.g. _datacopied and asarray):

CholeskySolve (potrs)
LUFactor (getrf)
SolveTriangular (trtrs)

Add coverage for empty case in each.

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--1605.org.readthedocs.build/en/1605/

codecov · 2025-09-02T16:17:39Z

Codecov Report

❌ Patch coverage is 60.86957% with 18 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.62%. Comparing base (46f8227) to head (1d3e180).
⚠️ Report is 43 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/tensor/slinalg.py	60.86%	9 Missing and 9 partials ⚠️

❌ Your patch status has failed because the patch coverage (60.86%) 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    #1605      +/-   ##
==========================================
- Coverage   81.64%   81.62%   -0.03%     
==========================================
  Files         231      231              
  Lines       52952    52992      +40     
  Branches     9388     9404      +16     
==========================================
+ Hits        43235    43257      +22     
- Misses       7273     7282       +9     
- Partials     2444     2453       +9

Files with missing lines	Coverage Δ
pytensor/tensor/slinalg.py	`91.71% <60.86%> (-1.70%)`	⬇️

🚀 New features to boost your workflow:

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

pytensor/tensor/slinalg.py

jessegrabowski

This is really great, thanks for taking this on!

I have one requested change (the one about shape checking) and the rest is related to what we should do if the algorithm fails. Please don't make those changes until the other devs weigh in, since it's a bit of an API break.

jessegrabowski · 2025-09-05T12:12:51Z

pytensor/tensor/slinalg.py

+        if c.ndim != 2 or c.shape[0] != c.shape[1]:
+            raise ValueError("The factored matrix c is not square.")
+        if c.shape[1] != b.shape[0]:
+            raise ValueError(f"incompatible dimensions ({c.shape} and {b.shape})")


We don't need to do shape checking in perform, that is handled by make_node

not true, shapes may not be static

not true, shapes may not be static

FWIW; we'll deprecate in-place modifications of the shape (also dtype and strides) modifications in numpy 2.4

Static here means we don't know the shape until runtime, as in the following graph:

import pytensor import pytensor.tensor as pt x = pt.vector("x", shape=(None,)) out = pt.exp(x) fn = pytensor.function([x], out) fn([1, 2, 3]) fn([1, 2, 3, 4]) # Allowed to call with different input lengths each time

jessegrabowski · 2025-09-05T12:14:12Z

pytensor/tensor/slinalg.py

+        if info != 0:
+            raise ValueError(f"illegal value in {-info}th argument of internal potrs")
+


I'd prefer if we returned a matrix of np.nan if info !=0 rather than erroring out. This is what jax does, and it makes it a lot more ergonomic to work with in iterative algorithms.

This might be out of scope for this PR; asking @ricardoV94 for a 2nd opinion

jessegrabowski · 2025-09-05T12:14:44Z

pytensor/tensor/slinalg.py

+        if info < 0:
+            raise ValueError(
+                f"illegal value in {-info}th argument of internal getrf (lu_factor)"
+            )
+        if info > 0:
+            warnings.warn(
+                f"Diagonal number {info} is exactly zero. Singular matrix.",
+                LinAlgWarning,
+                stacklevel=2,
+            )


jessegrabowski · 2025-09-05T12:15:01Z

pytensor/tensor/slinalg.py

+        if info > 0:
+            raise LinAlgError(
+                f"singular matrix: resolution failed at diagonal {info-1}"
+            )
+        elif info < 0:
+            raise ValueError(f"illegal value in {-info}-th argument of internal trtrs")


Fyrebright · 2025-09-27T17:51:29Z

I have one requested change (the one about shape checking) and the rest is related to what we should do if the algorithm fails. Please don't make those changes until the other devs weigh in, since it's a bit of an API break.

Just wanted to check if there was any further thoughts on this?

jessegrabowski · 2025-09-30T03:21:06Z

Just wanted to check if there was any further thoughts on this?

Thanks for poking us on this. I had a conversation with @ricardoV94 today about it, and asked him to weigh in here

ricardoV94 · 2025-09-30T10:32:31Z

We should check shape in the perform unless the code fails gracefully if the shape is wrong. If you get wrong results silently or segfaults otherwise, that's no good.

Returning nan is fine by me, but perhaps better addressed in a separate PR where we do that for all linalg Ops.

jessegrabowski

LGTM

There was a failed test, but it didn't look like it was related to your changes (it was something about MinMax on integers). I'm rerunning it to see if it magically goes away, otherwise will need to look more closely.

ricardoV94 · 2025-10-01T08:49:09Z

pytensor/tensor/slinalg.py

+        if self.check_finite and not (np.isfinite(c).all() and np.isfinite(b).all()):
+            raise ValueError("array must not contain infs or NaNs")
+
+        if c.ndim != 2 or c.shape[0] != c.shape[1]:


To be precise, we don't need to check ndims, but I guess it's fine.

pytensor/tensor/slinalg.py

Fyrebright · 2025-10-01T20:25:50Z

There was a failed test, but it didn't look like it was related to your changes (it was something about MinMax on integers). I'm rerunning it to see if it magically goes away, otherwise will need to look more closely.

I was unable to replicate it locally. It should have xfailed (fast compile = 0) if I understand correctly.

jessegrabowski · 2025-10-01T21:50:35Z

I was unable to replicate it locally. It should have xfailed (fast compile = 0) if I understand correctly.

Looks like it magically fixed itself :)

jessegrabowski · 2025-10-12T16:22:55Z

@ricardoV94 any objections to merging this

ricardoV94 · 2025-10-12T17:43:18Z

Nope

jessegrabowski · 2025-10-12T17:55:28Z

@Fyrebright Thanks for work on this!

…#1605) * Use lapack instead of `scipy_linalg.cho_solve` * Use lapack instead of `scipy_linalg.lu_factor` * Use lapack instead of `scipy_linalg.solve_triangular` * Add empty test for lu_factor * Tidy imports * remove ndim check

Fyrebright marked this pull request as ready for review September 2, 2025 16:24

jorenham reviewed Sep 5, 2025

View reviewed changes

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

jessegrabowski requested changes Sep 5, 2025

View reviewed changes

Fyrebright force-pushed the 1468-chosolve-lufact-triangularsolve branch from a536532 to cebc540 Compare September 30, 2025 14:48

Fyrebright added 4 commits September 30, 2025 10:54

Use lapack instead of scipy_linalg.cho_solve

63e2ccc

Use lapack instead of scipy_linalg.lu_factor

3a8a9a4

Use lapack instead of scipy_linalg.solve_triangular

0197de5

Add empty test for lu_factor

334a44e

Fyrebright force-pushed the 1468-chosolve-lufact-triangularsolve branch from cebc540 to 334a44e Compare September 30, 2025 14:57

Tidy imports

8061159

jessegrabowski approved these changes Oct 1, 2025

View reviewed changes

ricardoV94 reviewed Oct 1, 2025

View reviewed changes

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

remove ndim check

1d3e180

jessegrabowski merged commit f48068a into pymc-devs:main Oct 12, 2025
63 of 64 checks passed

		if info != 0:
		raise ValueError(f"illegal value in {-info}th argument of internal potrs")

Use LAPACK functions for cho_solve, lu_factor, solve_triangular #1605

Use LAPACK functions for cho_solve, lu_factor, solve_triangular #1605

Conversation

Fyrebright commented Sep 2, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue

Checklist

Type of change

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

Uh oh!

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

jessegrabowski Sep 5, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Sep 5, 2025

Choose a reason for hiding this comment

Uh oh!

jorenham Sep 5, 2025

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Sep 6, 2025 • 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!

jessegrabowski Sep 5, 2025

Choose a reason for hiding this comment

Uh oh!

jessegrabowski Sep 5, 2025

Choose a reason for hiding this comment

Uh oh!

jessegrabowski Sep 5, 2025

Choose a reason for hiding this comment

Uh oh!

Fyrebright commented Sep 27, 2025

Uh oh!

jessegrabowski commented Sep 30, 2025

Uh oh!

ricardoV94 commented Sep 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jessegrabowski 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!

ricardoV94 Oct 1, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

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

Uh oh!

Uh oh!

jessegrabowski commented Oct 1, 2025

Uh oh!

jessegrabowski commented Oct 12, 2025

Uh oh!

ricardoV94 commented Oct 12, 2025

Uh oh!

Uh oh!

jessegrabowski commented Oct 12, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Use LAPACK functions for `cho_solve`, `lu_factor`, `solve_triangular` #1605

Use LAPACK functions for `cho_solve`, `lu_factor`, `solve_triangular` #1605

Fyrebright commented Sep 2, 2025 •

edited by github-actions bot

Loading

codecov bot commented Sep 2, 2025 •

edited

Loading

ricardoV94 Sep 6, 2025 •

edited

Loading

ricardoV94 commented Sep 30, 2025 •

edited

Loading

jessegrabowski left a comment •

edited

Loading

Fyrebright commented Oct 1, 2025 •

edited

Loading