Rewrite scalar solve to division

[jessegrabowski]

Description

Small optimization that rewrites solve(a,b) -> b / a when the core shape of a is (1, 1). This avoids calling a LAPACK routine in a case where it's simply not necessary.

This came up in the gradient of minimize, for cases where the function being minimized has only one input. The L_op in that case requires a bunch of linear algebra, but it can all be rewritten away when we're just dealing with scalars.

I also tweaked rewrite_cholesky_diag_to_sqrt_diag to apply to the (1, 1) case -- this allows cho_solve to be rewritten to just b / a ** 2, without any linalg calls. The only bummer is that the sqr and sqrt don't cancel. I get this graph:

Here is the graph, including the new sqr/sqrt cancellation:

import pytensor.tensor as pt
import pytensor
x = pt.tensor('x', shape=(1, 1))
c_and_lower = pt.linalg.cholesky(x), True
b = pt.tensor('b', shape=(None,))

f = pytensor.function([x, b], pt.linalg.cho_solve(c_and_lower, b))

f.dprint()

Composite{(i1 / abs(i0))} [id A] 2
 ├─ Squeeze{axis=1} [id B] 1
 │  └─ x [id C]
 └─ DimShuffle{order=[x]} [id D] 0
    └─ b [id E]

Inner graphs:

Composite{(i1 / abs(i0))} [id A]
 ← true_div [id F] 'o0'
    ├─ i1 [id G]
    └─ Abs [id H]
       └─ i0 [id I]

Related Issue

• Closes #
• Related to #

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://pytensor--1453.org.readthedocs.build/en/1453/

[ricardoV94]

Re sqr(sqrt) we should perhaps canonicalize to power, and have the pow(pow(x, a1), a2) -> pow(x, a1*a2) rewrite.

However we may need to introduce some nan switches if x < 0, depending on the values of a (and or refuse to rewrite if either is unknown)

[jessegrabowski]

Oh I didn't even think about the domain issue. That's a good reason we don't do that rewrite.

[ricardoV94]

Oh I didn't even think about the domain issue. That's a good reason we don't do that rewrite.

We handle log(exp) and exp(log), fine, it's just that depending on the order the rewrite has a nan switch or not

[jessegrabowski]

Common complex L

[jessegrabowski]

We had a test that assumed Cholesky of a 1x1 matrix should not be rewritten to sqrt(diag(x)), which is obviously wrong. I'm not sure why we had that, but I added a check for that case to the cholesky diag rewrite, and removed the test.

I also added the sqrt(sqr(x)) rewrites, so I did some git history cleanup. This PR should be a rebase merge with two commits, one for each feature.

[Copilot]

Pull Request Overview

This PR introduces optimizations for simplifying scalar operations and small linear algebra expressions.

• Adds rewrites to cancel nested sqrt/sqr operations.
• Optimizes solve, solve_triangular, and cho_solve on 1×1 matrices to simple division.
• Updates tests to cover these new rewrites.

Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
tests/tensor/rewriting/test_math.py	Added TestSqrSqrt to verify sqrt/sqr cancellation rules.
tests/tensor/rewriting/test_linalg.py	Added test_scalar_solve_to_division_rewrite for 1×1 solves.
pytensor/tensor/rewriting/math.py	Defined local_sqrt_sqr rewrite to eliminate nested ops.
pytensor/tensor/rewriting/linalg.py	Added 1×1 cholesky → sqrt rewrite and scalar_solve_to_divison.

[codecov]

Codecov Report

Attention: Patch coverage is 89.36170% with 5 lines in your changes missing coverage. Please review.

Project coverage is 81.99%. Comparing base (236e50d) to head (8ebf636).

Files with missing lines	Patch %	Lines
pytensor/tensor/rewriting/linalg.py	88.46%	2 Missing and 1 partial ⚠️
pytensor/tensor/rewriting/math.py	90.47%	1 Missing and 1 partial ⚠️

❌ Your patch check has failed because the patch coverage (89.36%) 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    #1453   +/-   ##
=======================================
  Coverage   81.98%   81.99%           
=======================================
  Files         231      231           
  Lines       52192    52239   +47     
  Branches     9185     9197   +12     
=======================================
+ Hits        42790    42832   +42     
- Misses       7094     7097    +3     
- Partials     2308     2310    +2     

Files with missing lines	Coverage Δ
pytensor/tensor/rewriting/math.py	89.50% <90.47%> (+0.01%)	⬆️
pytensor/tensor/rewriting/linalg.py	92.04% <88.46%> (-0.24%)	⬇️

🚀 New features to boost your workflow:

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