Add rewrite to merge multiple SVD Ops with different settings #769

HangenYuu · 2024-05-14T13:41:33Z

Description

When there are two or more SVD Ops with the same inputs on a graph, differing only by compute_uv, compute_uv = False should be changed to True everywhere. This will allow pytensor to see that these outputs are equivalent and re-use them, rather than computing the decomposition multiple times.

Related Issue

Closes Add rewrite to merge multiple SVD Ops with different settings #732
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

…nsor into svd_graph_rewrite

HangenYuu · 2024-05-14T14:01:24Z

The PR is still draft right now. I have added a minimally modified copy of tensor\rewritings\linalg\local_det_chol to tensor\rewritings\linalg. I have the following questions:

Am I using the APIs correctly to access and/or modify the argument/attribute of an Op?
I have been tweaking a small example involving computing gradient s w.r.t input a to check for the effect of the rewrite:

import pytensor
import pytensor.tensor as pt
import numpy as np
from pytensor.tensor.type import matrix
from pytensor.tensor.linalg import svd

a_pt = matrix("a")
s = svd(a_pt, full_matrices=False, compute_uv=False)
J, updates = pytensor.scan(lambda i, s, a_pt : pt.grad(s[i], a_pt), sequences=pt.arange(s.shape[0]), non_sequences=[s, a_pt])
f = pytensor.function([a_pt], J, updates=updates)
e = pytensor.graph.fg.FunctionGraph([a_pt], [J], clone=False)

which produces a graph for f with 2 SVDs differing only compute_uv as required.

However, the graph after rewriting of e contains only 1 SVD so the effect is masked.

Tweaking either ended up in the same situation or led to TypeError: Cost must be a scalar. e.g., this Hessian example

import pytensor
import pytensor.tensor as pt
from pytensor.tensor.type import matrix
from pytensor.tensor.linalg import svd

a_pt = matrix("a")
s = svd(a_pt, full_matrices=False, compute_uv=False)
gy = pt.grad(pt.sum(s), a_pt)
H, updates = pytensor.scan(lambda i, gy, a_pt : pt.grad(gy[i], a_pt), sequences=pt.arange(gy.shape[0]), non_sequences=[gy, a_pt])
f = pytensor.function([a_pt], H, updates=updates)
e = pytensor.graph.fg.FunctionGraph([a_pt], [H], clone=False)

Do you have suggestion for a small example to test the rewrite? This one can later be reused for unit testing.

aseyboldt · 2024-05-14T14:12:23Z

pytensor/tensor/rewriting/linalg.py

+    if svd_count > 1 and compute_uv:
+        for cl in not_compute_uv_svd_list:
+            cl.op.core_op.compute_uv = True
+    return [cl.outputs[0] for cl in not_compute_uv_svd_list]


I think changing properties of the op inplace might lead to problems...

This rewrite function should run for each SVD node, so maybe it is easier to just locate an existing compute_uv = True node, and return that as replacement for each compuet_uv = False node?

So something like:

If compute_uv is False, return and do nothing

check if there is a compute_uv = True node in the graph with the same input. If not, return and do nothing

Return the exising output of that node as replacement for the current compute_uv = False node.

I wonder though if there could be bad interactions somewhere if there is a rewrite that replaces compute_uv = Fales nodes if they are not used? We don't want to run into any infinite cycles...

@ricardoV94 Do you know if there are any problems that could happen if a rewrite returns an existing variable instead of a new one?

I think there will be a problem only when a rewrite tries to replace a variable by another that depends on the original variable.

And yes we shouldn't modify the properties in place. We should replace the smaller Op by the bigger one, just make sure the smaller one is not in the ancestors of the bigger one.

Otherwise creating a new SVD should be simple, just call the user facing constructor with the specific flags

Sorry I seemed to dump information carelessly. The gist was

I updated the code logic to be a node rewriter.

The rewrite is registered properly in optdb. However, I am having trouble coming up with a test case to show the effect of the rewrite. Perhaps @jessegrabowski can provide the original use case that led to you opening the issue Add rewrite to merge multiple SVD Ops with different settings #732?

It will arise in gradient graphs. For example, you can just do:

X = pt.dmatrix('X') s = pt.linalg.svd(X, compute_uv=False) g = pt.grad(s.sum(), X)

The graph for g will re-compute the SVD of X during the backward pass with compute_uv = True, because we require the matrices U and V to compute the gradient of s with respect to X. Pytensor then won't be able to see that these two computations are the same, and will end up computing the SVD twice.

a_pt = matrix("a") s = svd(a_pt, full_matrices=False, compute_uv=False) gs = pt.grad(pt.sum(s), a_pt) f = pytensor.function([a_pt], gs) e = pytensor.graph.fg.FunctionGraph([a_pt], [gs], clone=False)

Thank you. I indeed received a graph for gs and e with 2 different SVD:

But for f, I receive a graph with just a single SVD (that seems to be rewritten already with compute_uv=True):

The f's rewritten graph will be used in calculation if I run f([[1, 2], [3, 4]]). Does this satisfy your end goal already?

This is f summary profile:

Function profiling ================== Message: /tmp/ipykernel_1282122/871230895.py:10 Time in 1 calls to Function.__call__: 3.448710e-02s Time in Function.vm.__call__: 0.03426380921155214s (99.353%) Time in thunks: 0.03424406051635742s (99.295%) Total compilation time: 4.109558e-02s Number of Apply nodes: 2 PyTensor rewrite time: 2.893809e-02s PyTensor validate time: 2.457825e-04s PyTensor Linker time (includes C, CUDA code generation/compiling): 0.00876139895990491s C-cache preloading 5.506449e-03s Import time 8.061258e-04s Node make_thunk time 1.967770e-03s Node Dot22(SVD{full_matrices=False, compute_uv=True}.0, SVD{full_matrices=False, compute_uv=True}.2) time 1.942240e-03s Node SVD{full_matrices=False, compute_uv=True}(a) time 1.436425e-05s Time in all call to pytensor.grad() 1.036228e-02s Time since pytensor import 2.774s Class --- <% time> <sum %> <apply time> <time per call> <type> <#call> <#apply> <Class name> 99.8% 99.8% 0.034s 3.42e-02s Py 1 1 pytensor.tensor.nlinalg.SVD 0.2% 100.0% 0.000s 6.60e-05s C 1 1 pytensor.tensor.blas.Dot22 ... (remaining 0 Classes account for 0.00%(0.00s) of the runtime) Ops --- <% time> <sum %> <apply time> <time per call> <type> <#call> <#apply> <Op name> 99.8% 99.8% 0.034s 3.42e-02s Py 1 1 SVD{full_matrices=False, compute_uv=True} 0.2% 100.0% 0.000s 6.60e-05s C 1 1 Dot22 ... (remaining 0 Ops account for 0.00%(0.00s) of the runtime) Apply ------ <% time> <sum %> <apply time> <time per call> <#call> <id> <Apply name> 99.8% 99.8% 0.034s 3.42e-02s 1 0 SVD{full_matrices=False, compute_uv=True}(a) 0.2% 100.0% 0.000s 6.60e-05s 1 1 Dot22(SVD{full_matrices=False, compute_uv=True}.0, SVD{full_matrices=False, compute_uv=True}.2) ... (remaining 0 Apply instances account for 0.00%(0.00s) of the runtime)

pytensor.dprint may be an easier way to introspect the graphs

codecov · 2024-05-18T13:11:37Z

Codecov Report

Attention: Patch coverage is 26.08696% with 17 lines in your changes are missing coverage. Please review.

Project coverage is 80.82%. Comparing base (8c157a2) to head (8ba5119).
Report is 36 commits behind head on main.

❗ Current head 8ba5119 differs from pull request most recent head 1c30ee9

Please upload reports for the commit 1c30ee9 to get more accurate results.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #769      +/-   ##
==========================================
- Coverage   80.85%   80.82%   -0.03%     
==========================================
  Files         162      162              
  Lines       47016    47067      +51     
  Branches    11501    11529      +28     
==========================================
+ Hits        38014    38044      +30     
- Misses       6751     6767      +16     
- Partials     2251     2256       +5

Files	Coverage Δ
pytensor/tensor/rewriting/linalg.py	`81.05% <26.08%> (-7.64%)`	⬇️

... and 6 files with indirect coverage changes

.gitignore

ricardoV94 · 2024-05-21T13:49:41Z

pytensor/tensor/rewriting/linalg.py

+    (x,) = node.inputs
+    compute_uv = False
+
+    for cl, _ in fgraph.clients[x]:


You have to be careful because if the output of the SVD is an output of the function one of the clients will be a string "output" and the call cl.op will fail.

ricardoV94 · 2024-05-21T13:50:32Z

pytensor/tensor/rewriting/linalg.py

+
+    for cl, _ in fgraph.clients[x]:
+        if isinstance(cl.op, Blockwise) and isinstance(cl.op.core_op, SVD):
+            if (not compute_uv) and cl.op.core_op.compute_uv:


I don't think you need that first check?

Suggested change

if (not compute_uv) and cl.op.core_op.compute_uv:

if cl.op.core_op.compute_uv:

ricardoV94 · 2024-05-21T13:51:32Z

pytensor/tensor/rewriting/linalg.py

+
+    for cl, _ in fgraph.clients[x]:
+        if isinstance(cl.op, Blockwise) and isinstance(cl.op.core_op, SVD):
+            if (not compute_uv) and cl.op.core_op.compute_uv:


You should check if the uv outputs of this node are actually used (i.e., they have clients of their own). If not, they are useless and the rewrite shouldn't happen. In fact, this or another rewrite should change the flag from True to False for those nodes

ricardoV94 · 2024-05-21T14:08:38Z

I would break this rewrite into different logical parts:

Find all SVD clients from the same input X
Check if any have compute_uv that is actually being used (has clients of their own).
If compute_uv is ever needed/used, replace any variable coming out of an SVD with compute_uv == False by one coming out of an SVD with compute_uv == True. You can return a dictionary of replacements {var_from_svd_without_uv: var_from_svd_with_uv, ...}. You should never have to create a new SVD for this case, because compute_uv can only ever be needed if at least one of the nodes already has it set to True and is using those variables elsewhere in the graph.
If compute_uv is never needed, replace any variable with compute_uv == True, by one of the existing ones with compute_uv==False. If there is no replacement, you can create a brand new SVD operation.

HangenYuu · 2024-05-22T07:52:56Z

pytensor/tensor/rewriting/linalg.py

+
+
+@register_canonicalize
+@register_stabilize
+@register_specialize
+@node_rewriter([SVD])
+def local_svd_uv_simplify(fgraph, node):
+    """If we have more than one `SVD` `Op`s and at least one has keyword argument
+    `compute_uv=True`, then we can change `compute_uv = False` to `True` everywhere
+    and allow `pytensor` to re-use the decomposition outputs instead of recomputing.
+    """
+    (x,) = node.inputs
+
+    if node.compute_uv:
+        # compute_uv=True returns [u, s, v].
+        # if at least u or v is used, no need to rewrite this node.
+        if (
+            fgraph.clients[node.outputs[0]] is not None
+            or fgraph.clients[node.outputs[2]] is not None
+        ):
+            return
+
+        # Else, has to replace the s of this node with s of an SVD Op that compute_uv=False.
+        # First, iterate to see if there is an SVD Op that can be reused.
+        for cl, _ in fgraph.clients[x]:
+            if cl == "output":
+                continue
+            if isinstance(cl.op, Blockwise) and isinstance(cl.op.core_op, SVD):
+                if not cl.op.core_op.compute_uv:
+                    return {fgraph.clients[node.outputs[1]]: cl.outputs[0]}
+
+        # If no SVD reusable, return a new one.
+        return [svd(x, full_matrices=node.full_matrices, compute_uv=False)]
+
+    else:
+        # compute_uv=False returns [s].
+        # We want rewrite if there is another one with compute_uv=True.
+        # For this case, just reuse the `s` from the one with compute_uv=True.
+        for cl, _ in fgraph.clients[x]:
+            if cl == "output":
+                continue
+            if isinstance(cl.op, Blockwise) and isinstance(cl.op.core_op, SVD):
+                if cl.op.core_op.compute_uv:
+                    return [cl.outputs[1]]


Thanks @ricardoV94. My understanding is like this: The SVD with compute_uv == False will return [s], while the one with compute_uv == True will return [u, s, v]. We want to rewrite when there are 2 SVD Ops using the same input in the graph with different compute_uv value. Let's take the specific example of 2 SVD Ops, svd_f which returns [s_f] and svd_t which returns [u_t, s_t, v_t]. Based on whether at least u_t or v_t is used (since we still have to calculate both even if we use just one of them for subsequent calculations), 1 of 2 rewrites can happen:

Case 1: If at least u_t or v_t is used: return [s_t] in place of [s_f].

Case 2: Else: return [s_f] in place of [s_t].

Case 3: Additionally, if there is just one SVD Op with compute_uv == True, but both u and v are not used, then it must be substituted with a new SVD Op with compute_uv == False.

Yup that's it!. When you write down the updated rewrite feel free to add comments with as much explanation as you did here!

There could also be some weird cases where there are 3 SVDs, one with uv and full_matrices that actually doesn't use the uv, and one with uv and not full matrices that actually uses them (or vice-versa). In that case we could replace one for the other, but perhaps that's too much to worry and unlikely to happen. I don't see we ignoring this causing any bug. I am just raising attention to it so we don't accidentally rewrite a full-matrices into non full-matrices that are actually used.

For this one return {fgraph.clients[node.outputs[1]]: cl.outputs[0]} is this the correct syntax?

Yup, that tells to replace the key by the value variable

pytensor/tensor/rewriting/linalg.py

ricardoV94 · 2024-05-23T10:16:11Z

pytensor/tensor/rewriting/linalg.py

+            if cl == "output":
+                continue
+            if isinstance(cl.op, Blockwise) and isinstance(cl.op.core_op, SVD):
+                if cl.op.core_op.compute_uv:


We only want to do this if that other node is actually using the UV. If not we would actually want to replace that node by this one

That would be taken care by the first half at that node turn. As this is a local rewrite applied to all SVD node, each node will have its turn.

Even if you don't want to handle that other node there's no reason to rewrite this node into it. In general it's better to do as few rewrites as possible as every time a rewrite succeeds all other candidate rewrites are rerun (until an Equilibrium is achieved and nothing changes anymore).

On second thought I like your eager approach better, it's not readable. Since SVDs are rare we don't need to over optimize

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

HangenYuu · 2024-05-25T02:46:40Z

The tests run successfully.

ricardoV94 · 2024-05-27T10:12:05Z

tests/tensor/rewriting/test_linalg.py

+    s_1 = svd(a, full_matrices=False, compute_uv=False)
+    _, s_2, _ = svd(a, full_matrices=False, compute_uv=True)
+    # full_matrices = True is not supported for grad of svd
+    gs = pt.grad(pt.sum(s_1), a)


Explain that grad introduces the svd with compute_uv=True.

Assert you have two SVDs in the original graph of gs, with and without compute_uv. This will make the test more readable

Hi @ricardoV94, this may be a stupid question, but how can I access the graph of gs?

tests/tensor/rewriting/test_linalg.py

HangenYuu · 2024-05-29T11:49:51Z

I will be slower for the next 2 weeks. I am house looking right now, which should be over by then. I don't expect it to resemble a wedding preparation like this, but it is what it is. For the changes you suggested @ricardoV94 I will edit them in a slot of free time tomorrow.

ricardoV94 · 2024-05-29T12:09:07Z

No worries and best of luck!

pytensor/tensor/rewriting/linalg.py

ricardoV94 · 2024-05-30T04:07:27Z

pytensor/tensor/rewriting/linalg.py

+                    return {fgraph.clients[node.outputs[1]]: cl.outputs[0]}
+
+        # If no SVD reusable, return a new one.
+        return {


Todo for myself, check if this remove is needed, and if so whether it's also needed in the return above

HangenYuu · 2024-06-02T12:56:28Z

Thanks @ricardoV94 for your patience.

Quick updates: I added your suggestions. The tests are not passed right now. I am looking at it. It seems that the rewrite does not happen for the second case

=================================== FAILURES ===================================
______________________________ test_svd_uv_merge _______________________________

    def test_svd_uv_merge():
        a = matrix("a")
        s_1 = svd(a, full_matrices=False, compute_uv=False)
        _, s_2, _ = svd(a, full_matrices=False, compute_uv=True)
        _, s_3, _ = svd(a, full_matrices=True, compute_uv=True)
        u_4, s_4, v_4 = svd(a, full_matrices=False, compute_uv=True)
        # `grad` will introduces an SVD Op with compute_uv=True
        # full_matrices = True is not supported for grad of svd
        gs = pt.grad(pt.sum(s_1), a)
    
        # 1. compute_uv=False needs rewriting with compute_uv=True
        f_1 = pytensor.function([a], gs)
        nodes = f_1.maker.fgraph.apply_nodes
        svd_counter = 0
        for node in nodes:
            if isinstance(node.op, SVD):
                assert node.op.compute_uv
                svd_counter += 1
        assert svd_counter == 1
    
        # 2. compute_uv=True needs rewriting with compute=False, reuse node
        f_2 = pytensor.function([a], [s_1, s_2])
        nodes = f_2.maker.fgraph.apply_nodes
        svd_counter = 0
        for node in nodes:
            if isinstance(node.op, SVD):
>               assert not node.op.compute_uv
E               assert not True
E                +  where True = SVD(full_matrices=False,compute_uv=True).compute_uv
E                +    where SVD(full_matrices=False,compute_uv=True) = SVD{full_matrices=False, compute_uv=True}(a).op

HangenYuu added 5 commits May 12, 2024 22:14

Started SVD rewrite

23ad3c4

Added rewrite to fix svd_graph_rewrite

d4a89ad

Started SVD rewrite

d4a2f2f

Added rewrite to fix svd_graph_rewrite

03c8e6f

Merge branch 'svd_graph_rewrite' of https://github.com/HangenYuu/pyte…

839f479

…nsor into svd_graph_rewrite

aseyboldt reviewed May 14, 2024

View reviewed changes

HangenYuu and others added 3 commits May 16, 2024 21:12

Fixed logic error of SVD node local rewrite: tried to rewrite globally

4a1be4b

Fixed typo error

b999d68

Merge branch 'pymc-devs:main' into svd_graph_rewrite

14f89f8

ricardoV94 reviewed May 21, 2024

View reviewed changes

HangenYuu added 2 commits May 22, 2024 15:39

Refactored logic for SVD to support all 3 cases

55ad931

Reverted .gitignore

0337e9d

HangenYuu force-pushed the svd_graph_rewrite branch from de46cff to 0337e9d Compare May 22, 2024 07:50

HangenYuu commented May 22, 2024

View reviewed changes

HangenYuu marked this pull request as ready for review May 23, 2024 01:50

ricardoV94 reviewed May 23, 2024

View reviewed changes

HangenYuu and others added 2 commits May 25, 2024 08:59

Update pytensor/tensor/rewriting/linalg.py

ecc62ae

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

Added unittest for SVD rewrite

8ba5119

ricardoV94 reviewed May 27, 2024

View reviewed changes

ricardoV94 reviewed May 30, 2024

View reviewed changes

Added test cases for SVD rewrite

1c30ee9

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add rewrite to merge multiple SVD Ops with different settings #769

Add rewrite to merge multiple SVD Ops with different settings #769

HangenYuu commented May 14, 2024

HangenYuu commented May 14, 2024 •

edited

aseyboldt May 14, 2024

aseyboldt May 14, 2024

ricardoV94 May 14, 2024

ricardoV94 May 14, 2024

ricardoV94 May 14, 2024

HangenYuu May 16, 2024 •

edited

jessegrabowski May 19, 2024 •

edited

HangenYuu May 19, 2024 •

edited

HangenYuu May 19, 2024 •

edited

ricardoV94 May 21, 2024

codecov bot commented May 18, 2024 •

edited

ricardoV94 May 21, 2024

ricardoV94 May 21, 2024

ricardoV94 May 21, 2024

ricardoV94 commented May 21, 2024 •

edited

HangenYuu May 22, 2024

ricardoV94 May 22, 2024 •

edited

ricardoV94 May 22, 2024 •

edited

HangenYuu May 23, 2024

ricardoV94 May 23, 2024

ricardoV94 May 23, 2024

HangenYuu May 25, 2024

ricardoV94 May 26, 2024

ricardoV94 May 30, 2024

HangenYuu commented May 25, 2024

ricardoV94 May 27, 2024

HangenYuu Jun 2, 2024

HangenYuu commented May 29, 2024

ricardoV94 commented May 29, 2024

ricardoV94 May 30, 2024

HangenYuu commented Jun 2, 2024

	if (not compute_uv) and cl.op.core_op.compute_uv:
	if cl.op.core_op.compute_uv:

Add rewrite to merge multiple SVD Ops with different settings #769

Are you sure you want to change the base?

Add rewrite to merge multiple SVD Ops with different settings #769

Conversation

HangenYuu commented May 14, 2024

Description

Related Issue

Checklist

Type of change

HangenYuu commented May 14, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

HangenYuu May 16, 2024 • edited

Choose a reason for hiding this comment

jessegrabowski May 19, 2024 • edited

Choose a reason for hiding this comment

HangenYuu May 19, 2024 • edited

Choose a reason for hiding this comment

HangenYuu May 19, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented May 18, 2024 • edited

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 commented May 21, 2024 • edited

Choose a reason for hiding this comment

ricardoV94 May 22, 2024 • edited

Choose a reason for hiding this comment

ricardoV94 May 22, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

HangenYuu commented May 25, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

HangenYuu commented May 29, 2024

ricardoV94 commented May 29, 2024

Choose a reason for hiding this comment

HangenYuu commented Jun 2, 2024

HangenYuu commented May 14, 2024 •

edited

HangenYuu May 16, 2024 •

edited

jessegrabowski May 19, 2024 •

edited

HangenYuu May 19, 2024 •

edited

HangenYuu May 19, 2024 •

edited

codecov bot commented May 18, 2024 •

edited

ricardoV94 commented May 21, 2024 •

edited

ricardoV94 May 22, 2024 •

edited

ricardoV94 May 22, 2024 •

edited