Rewrite scalar solve to division #1453

jessegrabowski · 2025-06-07T14:19:50Z

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:

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 / sqr(sqrt(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 / sqr(sqrt(i0)))} [id A]
 ← true_div [id F] 'o0'
    ├─ i1 [id G]
    └─ sqr [id H]
       └─ sqrt [id I]
          └─ i0 [id J]

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

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

Copilot

Pull Request Overview

This PR optimizes scalar linear solves by rewriting them to simple division (and square‐root for Cholesky) when the matrix is (1,1), avoiding expensive LAPACK calls.

Add a rewriter (scalar_solve_to_divison) to turn solve/solve_triangular/cho_solve on 1×1 into division.
Extend rewrite_cholesky_diag_to_sqrt_diag to map a (1,1) Cholesky to sqrt.
Add tests in test_linalg.py to verify these rewrites discard the original Solve* ops.

Reviewed Changes

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

File	Description
tests/tensor/rewriting/test_linalg.py	Added parametrized tests for scalar solve rewrites
pytensor/tensor/rewriting/linalg.py	Implemented two rewriters for scalar solve cases

Comments suppressed due to low confidence (3)

pytensor/tensor/rewriting/linalg.py:1035

[nitpick] The function name scalar_solve_to_divison contains a typo. Rename it to scalar_solve_to_division for consistency.

def scalar_solve_to_divison(fgraph, node):

pytensor/tensor/rewriting/linalg.py:1034

Using [Blockwise] here means the rewrite won’t match SolveBase nodes. Specify the correct op class (e.g., SolveBase) so this rule can fire.

@node_rewriter([Blockwise])

tests/tensor/rewriting/test_linalg.py:32

[nitpick] Indentation for this import is inconsistent with the surrounding lines; align it to match the existing style.

    CholeskySolve,

Copilot · 2025-06-07T14:21:52Z

pytensor/tensor/rewriting/linalg.py

@@ -908,6 +910,11 @@ def rewrite_cholesky_diag_to_sqrt_diag(fgraph, node):
        return None

    [input] = node.inputs
+
+    # Check if input is a (1, 1) matrix
+    if all(input.type.broadcastable[:-2]):


The check for a (1,1) matrix is incorrect: it tests all dimensions except the last two. It should verify the last two dims are broadcastable, e.g. all(input.type.broadcastable[-2:]).

Suggested change

if all(input.type.broadcastable[:-2]):

if all(input.type.broadcastable[-2:]):

Copilot seems right

pytensor/tensor/rewriting/linalg.py

ricardoV94 · 2025-06-07T15:19:12Z

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 · 2025-06-08T05:46:43Z

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

ricardoV94 · 2025-06-08T06:49:03Z

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 · 2025-06-08T07:04:40Z

Common complex L

ricardoV94 · 2025-06-07T15:14:26Z

pytensor/tensor/rewriting/linalg.py

+
+    # Special handling for different types of solve
+    match core_op:
+        case SolveTriangular():


This is expensive? Creating an Op in every match case?

Definitely. I was just being too infatuated with match.

Match doesn't work with the class (instead of instance) as an instance check?

No. See the match PEP

It sounds like it could be clever and not really instantiate the Op in the match check. Can you test with a small example (need not be PyTensor specific)

Yeah it doesn't instantiate stuff:

class AB: def __init__(self): print("AB was initialized") class AA(AB): def __init__(self): print("AA was initialized") class BB(AB): def __init__(self): print("BB was initialized") aa = AA() bb = BB() print("Starting match") for x in (aa, bb, 1): match x: case AA(): print("match aa") case BB(): print("match bb") case AB(): print("match ab") case _: print("match nothing") # AA was initialized # BB was initialized # Starting match # match aa # match bb # match nothing

Not sure how it's implemented, very surprising syntax imo, but seems fine to use like you did

https://stackoverflow.com/a/75749115

Wonder if we could start using it for match x.owner: Apply(op=Elemwise(scalar_op=Add)) kind of stuff instead of x.owner is not None and isinstance(x.owner.op, Elemwise) and isinstance(x.owner.op, Add)) 🤔

ricardoV94 · 2025-06-07T15:21:22Z

pytensor/tensor/rewriting/linalg.py

@@ -908,6 +910,11 @@ def rewrite_cholesky_diag_to_sqrt_diag(fgraph, node):
        return None

    [input] = node.inputs
+
+    # Check if input is a (1, 1) matrix
+    if all(input.type.broadcastable[:-2]):


Copilot seems right

jessegrabowski added 2 commits June 7, 2025 21:44

Rewrite scalar solve to division

f3b628d

Check for SolveBase to catch all cases

b9709e2

jessegrabowski requested review from ricardoV94 and Copilot June 7, 2025 14:19

Copilot AI reviewed Jun 7, 2025

View reviewed changes

jessegrabowski added enhancement New feature or request graph rewriting linalg Linear algebra labels Jun 7, 2025

ricardoV94 requested changes Jun 8, 2025

View reviewed changes

	if all(input.type.broadcastable[:-2]):
	if all(input.type.broadcastable[-2:]):

Rewrite scalar solve to division #1453

Are you sure you want to change the base?

Rewrite scalar solve to division #1453

Conversation

jessegrabowski commented Jun 7, 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!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Copilot AI Jun 7, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

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

Uh oh!

Uh oh!

jessegrabowski commented Jun 8, 2025

Uh oh!

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

Uh oh!

Uh oh!

jessegrabowski commented Jun 8, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jessegrabowski commented Jun 7, 2025 •

edited by github-actions bot

Loading

ricardoV94 commented Jun 7, 2025 •

edited

Loading

ricardoV94 commented Jun 8, 2025 •

edited

Loading