Add rewrite to lift linear algebra through certain linalg ops

[jessegrabowski]

Adds rewrites to lift certain linear algebra ops (inv, pinv, cholesky) through "compositions" of matrices (block_diag, kron).

Kron is currently broken, so it's marked as a draft.

Description

Related Issue

• Closes #
• Related to ENH: Implement COLA library rewrites for linear algebra functions #573

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):

[jessegrabowski]

I tracked down the bug we hit during the hackathon. We were importing pytensor.compile.builders.function, which is a module. We wanted pytensor.compile.builders.function.function.

Now the problem is that kron isn't well defined on batches. We noted that it works for tensors of any dimension, but it doesn't define batched behavior. If we have tensors of shape (batch, a, b) and (batch, c, d), the expected output of a batched kron is shape (batch, a*c, b*d), but we get something I really don't understand. Examples:

from scipy import linalg
import numpy as np
import pytensor.tensor as pt

rng = np.random.default_rng()

# Desired
A = rng.normal(size=(4, 2, 3))
B = rng.normal(size=(4, 5, 6))

x = np.vectorize(linalg.kron, signature='(n,m),(o,p)->(q,r)')(A, B)
x.shape  # Out: (4, 10, 18)

x = pt.linalg.kron(A, B)
x.eval() # Out: (12, 8, 5, 6)

So our batched tests for this rewrite currently failed for the batched kron. Since the base case is defined on arbitrary dimensions, it seems like our options are:

1. Have users define their own BlockWise kron as needed. Apply the test to a customized batched kron.
2. Add an axis argument to kron that lets the user specify which axes are the core dimensions. It could default to None (all dimensions are core), but we test on axis=[-1,-2].
3. xfail the kron batched test.

Thoughts?

[ricardoV94]

Call pt.vectorize on the test from the 2D core case?

Does the rewrite still hold as is though?

[jessegrabowski]

The rewrite works fine on the 2d case, it just fails on the batch case because the way kron currently works breaks the "liftability" if ndim > 2.

This is actually an artifact of how we define kron. It seems like there are two ways: np.kron is not the same as scipy.linalg,kron. Theano went with scipy, JAX went with numpy. The rewrite would work with the numpy definition:

np.allclose(np.kron(np.linalg.pinv(A), np.linalg.pinv(B)),
            np.linalg.pinv(np.kron(A, B)))
# True

[ricardoV94]

Okay, so maybe let's add a kwarg: scipy_like? I wouldn't add axis just because it's more complex for us. Numpy behavior is simpler in that axis are always the last 2, which we can get automatically with Blockwise?

Regardless, we have to make sure we do the rewrite correctly (or restrict it to the cases where the logic is correct, although it doesn't sound like it should be too complicated to support both forms?)

[jessegrabowski]

For the scipy form, we need a way for the user to tell us what are batch axes and what are not. If we don't know, the rewrite will break for anything bigger than 2d. I guess if we had a mode keyword, we could just check it and skip the rewrite if it's scipy.

For the numpy form, we technically still need to know this (because Blockwise(kron(A, B)) isn't the same as kron(A,B) for ndim > 2), but the rewrite still works in all cases. If were were in this case, we could just provide kron "as is" and demand users vectorize it themselves as needed, as jax does.

[ricardoV94]

For the scipy form, we need a way for the user to tell us what are batch axes and what are not. If we don't know, the rewrite will break for anything bigger than 2d. I guess if we had a mode keyword, we could just check it and skip the rewrite if it's scipy.

Isn't our current implementation already the Scipy kron. We don't need to do anything. A numpy kron is a Blockwised kron with 2d core inputs, which right now users would have to create manually. For this PR you can decide to support only Scipy kron or also a Blockwised Kron?

I don't see why the axis is needed. Scipy mixes everything if I understand correctly, and Numpy only the last two axes. From the signature / absense of the Blockwise (and its signature) we should know which case it is.

[jessegrabowski]

Numpy is not blockwise, it just multiplies all the axes together in a way that still respects commutativity with matrix operations. Here's a guide to all the outputs:

# 2d case -- both are the same
A = rng.normal(size=(2, 3))
B = rng.normal(size=(5, 6))
np.allclose(np.kron(A, B), linalg.kron(A, B))
# True

# nd case -- both are different. Note the effect on the "batch" dimension in the numpy case (becomes 4*4=16)
A = rng.normal(size=(4, 2, 3))
B = rng.normal(size=(4, 5, 6))

print(np.kron(A,B).shape)
print(linalg.kron(A,B).shape)

# (16, 10, 18)
# (12, 8, 5, 6)

# Vectorized 2d case -- both are the same again,  but not the same as the numpy nd case
vec_np_kron = np.vectorize(np.kron, signature='(n,m),(o,p)->(q,r)')
vec_sp_kron = np.vectorize(linalg.kron, signature='(n,m),(o,p)->(q,r)')

print(vec_np_kron(A,B).shape)
print(vec_sp_kron(A,B).shape)

# (4, 10, 18)
# (4, 10, 18)

[jessegrabowski]

After #684 this should be ready to go. I touched pytensor/link/c/cmodule.py because mypy was complaining about it on my end, but now mypy is complaining about something on the CI, so I'll need some guidance on what to do there.

[ricardoV94]

After #684 this should be ready to go. I touched pytensor/link/c/cmodule.py because mypy was complaining about it on my end, but now mypy is complaining about something on the CI, so I'll need some guidance on what to do there.

Things like this have happened in the past. My best hunch is a difference between individuals commits vs all files. Locally is doing one and in CI the other and mypy opinion changes accordingly.

Or python version.

Or something else

[jessegrabowski]

Probably the eclipse for all I know. Mypy is something else.

[ricardoV94]

Nice! Proved trickier because of the scipy thing but not terrible in the end

[ricardoV94]

Can you add a new test for the error and leave this one unchanged? Or did this start failing? Looks like it shouldn't have?

[jessegrabowski]

Yes, it started failing.

[ricardoV94]

Hmm why? It was working before, may need to investigate or is it obvious?

[jessegrabowski]

It merits investigation. But it might be that the test was always flakey and now we're catching it. The duplicated input is one of the intermediate f representations -- something like *4 Matrix<?,?>. The true inputs x,y,z aren't the problem. I assume this is something to do with how the grad function works?

[ricardoV94]

Intermediate variables can be repeated just fine

[jessegrabowski]

The test doesn't have obvious double inputs though, what's going on there? Is it because the original f is considered input to both grad(f.sum(), y) and grad(grad(f.sum(), y).sum(), y) ?

[ricardoV94]

Let me see what are the duplicate inputs. In any case you proved we are too restrictive with the OFG([a], [a]) example.

[ricardoV94]

The L_op of the first grad call looks like:

pytensor.dprint(self.fgraph, print_fgraph_inputs=True)

→ *0-<Matrix(float64, shape=(?, ?))> [id A]
→ *1-<Matrix(float64, shape=(?, ?))> [id B]
→ *2-<Matrix(float64, shape=(?, ?))> [id C]
→ *3-<Matrix(float64, shape=(?, ?))> [id D]
→ *4-<Matrix(float64, shape=(?, ?))> [id E]
*4-<Matrix(float64, shape=(?, ?))> [id E]
Mul [id F] 0
 ├─ *4-<Matrix(float64, shape=(?, ?))> [id E]
 └─ *2-<Matrix(float64, shape=(?, ?))> [id C]
Mul [id G] 1
 ├─ *4-<Matrix(float64, shape=(?, ?))> [id E]
 └─ *1-<Matrix(float64, shape=(?, ?))> [id B]

You can see the last input 4* (the output_grad) is also the first output (the input grad wrt to x is just the output grad). So when the second L_op is called on the first one, it will try to build an OFG that takes as inputs 4* twice, once from its inputs, and once from its outputs.

[jessegrabowski]

I went with skipping the input check if on_unused_inputs == 'ignore'. I thought that was better than removing the error. I guess a real solution will be to check for outputs that are inputs when we build the OpFromGraph, and clone them if so?

[ricardoV94]

Agree, but let's do that in a separate PR. Then we can get rid of this hacky params as well

[jessegrabowski]

float32 😍

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 80.78%. Comparing base (14651fb) to head (c1da527).
Report is 7 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #622      +/-   ##
==========================================
+ Coverage   80.76%   80.78%   +0.02%     
==========================================
  Files         162      162              
  Lines       46707    46757      +50     
  Branches    11422    11440      +18     
==========================================
+ Hits        37723    37773      +50     
  Misses       6732     6732              
  Partials     2252     2252              

Files	Coverage Δ
pytensor/compile/builders.py	77.45% <100.00%> (+0.27%)	⬆️
pytensor/tensor/nlinalg.py	94.76% <100.00%> (+0.50%)	⬆️
pytensor/tensor/rewriting/linalg.py	88.69% <100.00%> (+1.70%)	⬆️

... and 5 files with indirect coverage changes

[jessegrabowski]

I squashed this PR down to two commits, one for the actual rewrite and one for the changes that touched OpFromGraph (that hopefully we can revert down the road). After the CI runs again I'll merge. The codecov is failing because we don't have tests for the OpFromGraph error -- do you care? I can handle this in the next PR that will deal less hackishly with the duplicated inputs/outputs problem.

[ricardoV94]

Yeah we can merge as is

[ricardoV94]

I wouldn't remove the check later though, just not use kwargs. So a test is not a bad idea, just don't want to drag the PR too much more

[ricardoV94]

Looks good