Use lapack func instead of scipy.linalg.cholesky

[Fyrebright]

• 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/

[ricardoV94]

Nice start, left some comments

[Fyrebright]

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

[jessegrabowski]

Suggested change

	eye = np.eye(1, dtype=x.dtype)
	(potrf,) = scipy_linalg.get_lapack_funcs(("potrf",), (eye,))
	c, _ = potrf(eye, lower=False, overwrite_a=False, clean=True)
	out[0] = np.empty_like(x, dtype=c.dtype)
	(potrf,) = scipy_linalg.get_lapack_funcs(("potrf",), (x,))
	out[0] = np.empty_like(x, dtype=potrf.dtype)

Since you just want the dtype, you can check potrf.dtype, you don't need to actually call the routine. It's also fine to pass in the empty x in this case, since it's only used for its dtype attribute (the data or lack thereof doesn't matter)

[jessegrabowski]

I don't like the introduction of a new variable. We know for sure x is a numpy array inside this function, so the only relevant code in np.asarray_chkfinite(x) is the if not np.isfinite(x).all(): part. We should just directly do that

[jessegrabowski]

Get potrf once before this size check and re-use it in all branches

[jessegrabowski]

save a variable c_contiguous_input and use that for checks later (instead of maintaining two arrays)