Rewrite solves involving kron to eliminate kron

[jessegrabowski]

Description

Rewrite graphs of the form solve(kron(A, B), x) to solve(A, solve(B, x.reshape).mT).mT.reshape. This eliminates the kronecker product, and provides significant speedup.

Important limitation is that it only covers the case when b_ndim=1, because the math underpinning the rewrite requires that x is a vector. This is still an important case, however, because it's what arises in the logp of a multivariate normal when the covariance matrix is kronecker.

Also I hit what appears to be a numerical bug in the batch case when assume_a = 'pos'. There is disagreement, but only in the 2nd row of the outputs. No matter the batch size, it's always the 2nd batch that has a numerical problem -- all other batches agree. I've left in the failing test for now. We don't even vectorize kron by default, so if I can't figure it out I might just disable the rewrite for the Blockwise(Kron) case for now.

Benchmarks follow, with:

• small: A, B are (10, 10)
• medium: A, B are (50, 50)
• large: A, B are (100, 100)

-----------------------------------------------------------------------------------------
Name (time in us)                                                            Min                       Max                      Mean                 StdDev                    Median                    IQR            Outliers          OPS            Rounds  Iterations
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
test_rewrite_solve_kron_to_solve_benchmark[no_rewrite-small]             17.2500 (1.0)             34.6250 (1.0)             18.5068 (1.0)           2.3577 (1.0)             17.6670 (1.0)           0.8755 (1.0)          8;10  54,034.2951 (1.0)          93           1
test_rewrite_solve_kron_to_solve_benchmark[rewrite-small]                19.2910 (1.12)            98.7500 (2.85)            21.9831 (1.19)          4.0015 (1.70)            20.9160 (1.18)          3.6250 (4.14)       135;35  45,489.5626 (0.84)       3261           1

test_rewrite_solve_kron_to_solve_benchmark[no_rewrite-medium]        93,532.8330 (>1000.0)     96,359.5420 (>1000.0)     94,835.3042 (>1000.0)     857.3874 (363.65)      94,672.9585 (>1000.0)   1,327.0000 (>1000.0)       3;0      10.5446 (0.00)         10           1
test_rewrite_solve_kron_to_solve_benchmark[rewrite-medium]               66.1660 (3.84)           288.5420 (8.33)            74.0905 (4.00)          8.2418 (3.50)            72.8750 (4.12)          5.9580 (6.81)      405;317  13,497.0108 (0.25)       7247           1

test_rewrite_solve_kron_to_solve_benchmark[no_rewrite-large]      3,250,903.0000 (>1000.0)  3,333,840.0830 (>1000.0)  3,300,615.3582 (>1000.0)  31,539.6476 (>1000.0)  3,302,135.1670 (>1000.0)  38,780.1145 (>1000.0)       2;0       0.3030 (0.00)          5           1
test_rewrite_solve_kron_to_solve_benchmark[rewrite-large]               183.1670 (10.62)          357.8750 (10.34)          196.7968 (10.63)        11.5470 (4.90)           194.6250 (11.02)         8.7920 (10.04)     401;206   5,081.3837 (0.09)       3442           1
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

Related Issue

• Closes Rewrite Solve involving Kron #1557
• 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--1559.org.readthedocs.build/en/1559/

[Copilot]

Pull Request Overview

This PR introduces a rewrite optimization that eliminates Kronecker products from linear solve operations by transforming solve(kron(A, B), x) to an equivalent expression using nested solves. This provides significant performance improvements, especially for larger matrices, by avoiding the computational cost of forming the Kronecker product.

Key changes:

• Adds rewrite_solve_kron_to_solve function that rewrites solve operations involving Kronecker products
• Includes comprehensive tests covering different matrix shapes, solve types, and batch operations
• Handles both regular and batched Kronecker products through Blockwise operations

Reviewed Changes

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

File	Description
pytensor/tensor/rewriting/linalg.py	Implements the core rewrite logic to transform solve(kron(A,B), x) to nested solve operations
tests/tensor/rewriting/test_linalg.py	Adds comprehensive test coverage for the new rewrite functionality

Comments suppressed due to low confidence (1)

pytensor/tensor/rewriting/linalg.py:896

• Missing parentheses around the condition. The line should be (isinstance(A.owner.op, Blockwise) to properly group with the and operator on the next line.

        or isinstance(A.owner.op, Blockwise)

[ricardoV94]

b_ndim=2 is just a batched vector solve, could you move it to a batch dimension of b?

import numpy as np
import pytensor.tensor as pt

a = pt.matrix("a", shape=(3, 3))
b = pt.matrix("b", shape=(3, 5))
out1 = pt.linalg.solve(a, b, b_ndim=2)
out2 = pt.moveaxis(pt.linalg.solve(a, pt.moveaxis(b, -1, 0), b_ndim=1), 0, -1)

rng = np.random.default_rng(20)
test_dict = {
    a: rng.normal(size=(a.type.shape)),
    b: rng.normal(size=(b.type.shape))
}
np.testing.assert_allclose(out1.eval(test_dict), out2.eval(test_dict))