diff --git a/pytensor/compile/mode.py b/pytensor/compile/mode.py index f80dfaaf5c..63a1ba835b 100644 --- a/pytensor/compile/mode.py +++ b/pytensor/compile/mode.py @@ -490,6 +490,8 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs): "fusion", "inplace", "scan_save_mem_prealloc", + "reuse_lu_decomposition_multiple_solves", + "scan_split_non_sequence_lu_decomposition_solve", ], ), ) diff --git a/pytensor/scan/rewriting.py b/pytensor/scan/rewriting.py index b8e6b009d8..c49fbadce4 100644 --- a/pytensor/scan/rewriting.py +++ b/pytensor/scan/rewriting.py @@ -2561,7 +2561,6 @@ def scan_push_out_dot1(fgraph, node): position=1, ) - scan_seqopt1.register( "scan_push_out_non_seq", in2out(scan_push_out_non_seq, ignore_newtrees=True), @@ -2569,10 +2568,9 @@ def scan_push_out_dot1(fgraph, node): "fast_run", "scan", "scan_pushout", - position=2, + position=3, ) - scan_seqopt1.register( "scan_push_out_seq", in2out(scan_push_out_seq, ignore_newtrees=True), @@ -2580,7 +2578,7 @@ def scan_push_out_dot1(fgraph, node): "fast_run", "scan", "scan_pushout", - position=3, + position=4, ) @@ -2592,7 +2590,7 @@ def scan_push_out_dot1(fgraph, node): "more_mem", "scan", "scan_pushout", - position=4, + position=5, ) @@ -2605,7 +2603,7 @@ def scan_push_out_dot1(fgraph, node): "more_mem", "scan", "scan_pushout", - position=5, + position=6, ) scan_eqopt2.register( diff --git a/pytensor/tensor/__init__.py b/pytensor/tensor/__init__.py index c6b421d003..ce590f8228 100644 --- a/pytensor/tensor/__init__.py +++ b/pytensor/tensor/__init__.py @@ -114,6 +114,7 @@ def _get_vector_length_Constant(op: Op | Variable, var: Constant) -> int: # isort: off +import pytensor.tensor._linalg from pytensor.tensor import linalg from pytensor.tensor import special from pytensor.tensor import signal diff --git a/pytensor/tensor/_linalg/__init__.py b/pytensor/tensor/_linalg/__init__.py new file mode 100644 index 0000000000..767374b10b --- /dev/null +++ b/pytensor/tensor/_linalg/__init__.py @@ -0,0 +1,2 @@ +# Register rewrites +import pytensor.tensor._linalg.solve diff --git a/pytensor/tensor/_linalg/solve/__init__.py b/pytensor/tensor/_linalg/solve/__init__.py new file mode 100644 index 0000000000..1d85f4a66b --- /dev/null +++ b/pytensor/tensor/_linalg/solve/__init__.py @@ -0,0 +1,2 @@ +# Register rewrites in the database +import pytensor.tensor._linalg.solve.rewriting diff --git a/pytensor/tensor/_linalg/solve/rewriting.py b/pytensor/tensor/_linalg/solve/rewriting.py new file mode 100644 index 0000000000..ff1c74cdec --- /dev/null +++ b/pytensor/tensor/_linalg/solve/rewriting.py @@ -0,0 +1,198 @@ +from collections.abc import Container +from copy import copy + +from pytensor.graph import Constant, graph_inputs +from pytensor.graph.rewriting.basic import copy_stack_trace, in2out, node_rewriter +from pytensor.scan.op import Scan +from pytensor.scan.rewriting import scan_seqopt1 +from pytensor.tensor.basic import atleast_Nd +from pytensor.tensor.blockwise import Blockwise +from pytensor.tensor.elemwise import DimShuffle +from pytensor.tensor.rewriting.basic import register_specialize +from pytensor.tensor.rewriting.linalg import is_matrix_transpose +from pytensor.tensor.slinalg import Solve, lu_factor, lu_solve +from pytensor.tensor.variable import TensorVariable + + +def decompose_A(A, assume_a): + if assume_a == "gen": + return lu_factor(A, check_finite=False) + else: + raise NotImplementedError + + +def solve_lu_decomposed_system(A_decomp, b, b_ndim, assume_a, transposed=False): + if assume_a == "gen": + return lu_solve(A_decomp, b, b_ndim=b_ndim, trans=transposed) + else: + raise NotImplementedError + + +def _split_lu_solve_steps( + fgraph, node, *, eager: bool, allowed_assume_a: Container[str] +): + if not isinstance(node.op.core_op, Solve): + return None + + def get_root_A(a: TensorVariable) -> tuple[TensorVariable, bool]: + # Find the root variable of the first input to Solve + # If `a` is a left expand_dims or matrix transpose (DimShuffle variants), + # the root variable is the pre-DimShuffled input. + # Otherwise, `a` is considered the root variable. + # We also return whether the root `a` is transposed. + transposed = False + if a.owner is not None and isinstance(a.owner.op, DimShuffle): + if a.owner.op.is_left_expand_dims: + [a] = a.owner.inputs + elif is_matrix_transpose(a): + [a] = a.owner.inputs + transposed = True + return a, transposed + + def find_solve_clients(var, assume_a): + clients = [] + for cl, idx in fgraph.clients[var]: + if ( + idx == 0 + and isinstance(cl.op, Blockwise) + and isinstance(cl.op.core_op, Solve) + and (cl.op.core_op.assume_a == assume_a) + ): + clients.append(cl) + elif isinstance(cl.op, DimShuffle) and cl.op.is_left_expand_dims: + # If it's a left expand_dims, recurse on the output + clients.extend(find_solve_clients(cl.outputs[0], assume_a)) + return clients + + assume_a = node.op.core_op.assume_a + + if assume_a not in allowed_assume_a: + return None + + A, _ = get_root_A(node.inputs[0]) + + # Find Solve using A (or left expand_dims of A) + # TODO: We could handle arbitrary shuffle of the batch dimensions, just need to propagate + # that to the A_decomp outputs + A_solve_clients_and_transpose = [ + (client, False) for client in find_solve_clients(A, assume_a) + ] + + # Find Solves using A.T + for cl, _ in fgraph.clients[A]: + if isinstance(cl.op, DimShuffle) and is_matrix_transpose(cl.out): + A_T = cl.out + A_solve_clients_and_transpose.extend( + (client, True) for client in find_solve_clients(A_T, assume_a) + ) + + if not eager and len(A_solve_clients_and_transpose) == 1: + # If theres' a single use don't do it... unless it's being broadcast in a Blockwise (or we're eager) + # That's a "reuse" inside the inner vectorized loop + batch_ndim = node.op.batch_ndim(node) + (client, _) = A_solve_clients_and_transpose[0] + original_A, b = client.inputs + if not any( + a_bcast and not b_bcast + for a_bcast, b_bcast in zip( + original_A.type.broadcastable[:batch_ndim], + b.type.broadcastable[:batch_ndim], + strict=True, + ) + ): + return None + + A_decomp = decompose_A(A, assume_a=assume_a) + + replacements = {} + for client, transposed in A_solve_clients_and_transpose: + _, b = client.inputs + b_ndim = client.op.core_op.b_ndim + new_x = solve_lu_decomposed_system( + A_decomp, b, b_ndim=b_ndim, assume_a=assume_a, transposed=transposed + ) + [old_x] = client.outputs + new_x = atleast_Nd(new_x, n=old_x.type.ndim).astype(old_x.type.dtype) + copy_stack_trace(old_x, new_x) + replacements[old_x] = new_x + + return replacements + + +def _scan_split_non_sequence_lu_decomposition_solve( + fgraph, node, *, allowed_assume_a: Container[str] +): + """If the A of a Solve within a Scan is a function of non-sequences, split the LU decomposition step. + + The LU decomposition step can then be pushed out of the inner loop by the `scan_pushout_non_sequences` rewrite. + """ + scan_op: Scan = node.op + non_sequences = set(scan_op.inner_non_seqs(scan_op.inner_inputs)) + new_scan_fgraph = scan_op.fgraph + + changed = False + while True: + for inner_node in new_scan_fgraph.toposort(): + if ( + isinstance(inner_node.op, Blockwise) + and isinstance(inner_node.op.core_op, Solve) + and inner_node.op.core_op.assume_a in allowed_assume_a + ): + A, b = inner_node.inputs + if all( + (isinstance(root_inp, Constant) or (root_inp in non_sequences)) + for root_inp in graph_inputs([A]) + ): + if new_scan_fgraph is scan_op.fgraph: + # Clone the first time to avoid mutating the original fgraph + new_scan_fgraph, equiv = new_scan_fgraph.clone_get_equiv() + non_sequences = {equiv[non_seq] for non_seq in non_sequences} + inner_node = equiv[inner_node] # type: ignore + + replace_dict = _split_lu_solve_steps( + new_scan_fgraph, + inner_node, + eager=True, + allowed_assume_a=allowed_assume_a, + ) + assert ( + isinstance(replace_dict, dict) and len(replace_dict) > 0 + ), "Rewrite failed" + new_scan_fgraph.replace_all(replace_dict.items()) + changed = True + break # Break to start over with a fresh toposort + else: # no_break + break # Nothing else changed + + if not changed: + return + + # Return a new scan to indicate that a rewrite was done + new_scan_op = copy(scan_op) + new_scan_op.fgraph = new_scan_fgraph + new_outs = new_scan_op.make_node(*node.inputs).outputs + copy_stack_trace(node.outputs, new_outs) + return new_outs + + +@register_specialize +@node_rewriter([Blockwise]) +def reuse_lu_decomposition_multiple_solves(fgraph, node): + return _split_lu_solve_steps(fgraph, node, eager=False, allowed_assume_a={"gen"}) + + +@node_rewriter([Scan]) +def scan_split_non_sequence_lu_decomposition_solve(fgraph, node): + return _scan_split_non_sequence_lu_decomposition_solve( + fgraph, node, allowed_assume_a={"gen"} + ) + + +scan_seqopt1.register( + "scan_split_non_sequence_lu_decomposition_solve", + in2out(scan_split_non_sequence_lu_decomposition_solve, ignore_newtrees=True), + "fast_run", + "scan", + "scan_pushout", + position=2, +) diff --git a/pytensor/tensor/rewriting/linalg.py b/pytensor/tensor/rewriting/linalg.py index cd202fe3ed..af42bee236 100644 --- a/pytensor/tensor/rewriting/linalg.py +++ b/pytensor/tensor/rewriting/linalg.py @@ -75,6 +75,13 @@ def is_matrix_transpose(x: TensorVariable) -> bool: if ndims < 2: return False transpose_order = (*range(ndims - 2), ndims - 1, ndims - 2) + + # Allow expand_dims on the left of the transpose + if (diff := len(transpose_order) - len(node.op.new_order)) > 0: + transpose_order = ( + *(["x"] * diff), + *transpose_order, + ) return node.op.new_order == transpose_order return False diff --git a/tests/tensor/linalg/__init__.py b/tests/tensor/linalg/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/tensor/linalg/test_rewriting.py b/tests/tensor/linalg/test_rewriting.py new file mode 100644 index 0000000000..6f04fac5fb --- /dev/null +++ b/tests/tensor/linalg/test_rewriting.py @@ -0,0 +1,163 @@ +import numpy as np +import pytest + +from pytensor import config, function, scan +from pytensor.compile.mode import get_default_mode +from pytensor.gradient import grad +from pytensor.scan.op import Scan +from pytensor.tensor._linalg.solve.rewriting import ( + reuse_lu_decomposition_multiple_solves, + scan_split_non_sequence_lu_decomposition_solve, +) +from pytensor.tensor.blockwise import Blockwise +from pytensor.tensor.linalg import solve +from pytensor.tensor.slinalg import LUFactor, Solve, SolveTriangular +from pytensor.tensor.type import tensor + + +def count_vanilla_solve_nodes(nodes) -> int: + return sum( + ( + isinstance(node.op, Solve) + or (isinstance(node.op, Blockwise) and isinstance(node.op.core_op, Solve)) + ) + for node in nodes + ) + + +def count_lu_decom_nodes(nodes) -> int: + return sum( + ( + isinstance(node.op, LUFactor) + or ( + isinstance(node.op, Blockwise) and isinstance(node.op.core_op, LUFactor) + ) + ) + for node in nodes + ) + + +def count_lu_solve_nodes(nodes) -> int: + count = sum( + ( + isinstance(node.op, SolveTriangular) + or ( + isinstance(node.op, Blockwise) + and isinstance(node.op.core_op, SolveTriangular) + ) + ) + for node in nodes + ) + # Each LU solve uses two Triangular solves + return count // 2 + + +@pytest.mark.parametrize("transposed", (False, True)) +def test_lu_decomposition_reused_forward_and_gradient(transposed): + rewrite_name = reuse_lu_decomposition_multiple_solves.__name__ + mode = get_default_mode() + + A = tensor("A", shape=(2, 2)) + b = tensor("b", shape=(2, 3)) + + x = solve(A, b, assume_a="gen", transposed=transposed) + grad_x_wrt_A = grad(x.sum(), A) + fn_no_opt = function([A, b], [x, grad_x_wrt_A], mode=mode.excluding(rewrite_name)) + no_opt_nodes = fn_no_opt.maker.fgraph.apply_nodes + assert count_vanilla_solve_nodes(no_opt_nodes) == 2 + assert count_lu_decom_nodes(no_opt_nodes) == 0 + assert count_lu_solve_nodes(no_opt_nodes) == 0 + + fn_opt = function([A, b], [x, grad_x_wrt_A], mode=mode.including(rewrite_name)) + opt_nodes = fn_opt.maker.fgraph.apply_nodes + assert count_vanilla_solve_nodes(opt_nodes) == 0 + assert count_lu_decom_nodes(opt_nodes) == 1 + assert count_lu_solve_nodes(opt_nodes) == 2 + + # Make sure results are correct + rng = np.random.default_rng(31) + A_test = rng.random(A.type.shape, dtype=A.type.dtype) + b_test = rng.random(b.type.shape, dtype=b.type.dtype) + resx0, resg0 = fn_no_opt(A_test, b_test) + resx1, resg1 = fn_opt(A_test, b_test) + rtol = 1e-7 if config.floatX == "float64" else 1e-6 + np.testing.assert_allclose(resx0, resx1, rtol=rtol) + np.testing.assert_allclose(resg0, resg1, rtol=rtol) + + +@pytest.mark.parametrize("transposed", (False, True)) +def test_lu_decomposition_reused_blockwise(transposed): + rewrite_name = reuse_lu_decomposition_multiple_solves.__name__ + mode = get_default_mode() + + A = tensor("A", shape=(2, 2)) + b = tensor("b", shape=(2, 2, 3)) + + x = solve(A, b, transposed=transposed) + fn_no_opt = function([A, b], [x], mode=mode.excluding(rewrite_name)) + no_opt_nodes = fn_no_opt.maker.fgraph.apply_nodes + assert count_vanilla_solve_nodes(no_opt_nodes) == 1 + assert count_lu_decom_nodes(no_opt_nodes) == 0 + assert count_lu_solve_nodes(no_opt_nodes) == 0 + + fn_opt = function([A, b], [x], mode=mode.including(rewrite_name)) + opt_nodes = fn_opt.maker.fgraph.apply_nodes + assert count_vanilla_solve_nodes(opt_nodes) == 0 + assert count_lu_decom_nodes(opt_nodes) == 1 + assert count_lu_solve_nodes(opt_nodes) == 1 + + # Make sure results are correct + rng = np.random.default_rng(31) + A_test = rng.random(A.type.shape, dtype=A.type.dtype) + b_test = rng.random(b.type.shape, dtype=b.type.dtype) + resx0 = fn_no_opt(A_test, b_test) + resx1 = fn_opt(A_test, b_test) + np.testing.assert_allclose(resx0, resx1) + + +@pytest.mark.parametrize("transposed", (False, True)) +def test_lu_decomposition_reused_scan(transposed): + rewrite_name = scan_split_non_sequence_lu_decomposition_solve.__name__ + mode = get_default_mode() + + A = tensor("A", shape=(2, 2)) + x0 = tensor("b", shape=(2, 3)) + + xs, _ = scan( + lambda xtm1, A: solve(A, xtm1, assume_a="general", transposed=transposed), + outputs_info=[x0], + non_sequences=[A], + n_steps=10, + ) + + fn_no_opt = function( + [A, x0], + [xs], + mode=mode.excluding(rewrite_name), + ) + [no_opt_scan_node] = [ + node for node in fn_no_opt.maker.fgraph.apply_nodes if isinstance(node.op, Scan) + ] + no_opt_nodes = no_opt_scan_node.op.fgraph.apply_nodes + assert count_vanilla_solve_nodes(no_opt_nodes) == 1 + assert count_lu_decom_nodes(no_opt_nodes) == 0 + assert count_lu_solve_nodes(no_opt_nodes) == 0 + + fn_opt = function([A, x0], [xs], mode=mode.including("scan", rewrite_name)) + [opt_scan_node] = [ + node for node in fn_opt.maker.fgraph.apply_nodes if isinstance(node.op, Scan) + ] + opt_nodes = opt_scan_node.op.fgraph.apply_nodes + assert count_vanilla_solve_nodes(opt_nodes) == 0 + # The LU decomp is outside of the scan! + assert count_lu_decom_nodes(opt_nodes) == 0 + assert count_lu_solve_nodes(opt_nodes) == 1 + + # Make sure results are correct + rng = np.random.default_rng(170) + A_test = rng.random(A.type.shape, dtype=A.type.dtype) + x0_test = rng.random(x0.type.shape, dtype=x0.type.dtype) + resx0 = fn_no_opt(A_test, x0_test) + resx1 = fn_opt(A_test, x0_test) + rtol = 1e-7 if config.floatX == "float64" else 1e-6 + np.testing.assert_allclose(resx0, resx1, rtol=rtol) diff --git a/tests/tensor/test_blockwise.py b/tests/tensor/test_blockwise.py index a140e07846..cbaf27da29 100644 --- a/tests/tensor/test_blockwise.py +++ b/tests/tensor/test_blockwise.py @@ -328,7 +328,7 @@ class BlockwiseOpTester: @classmethod def setup_class(cls): - seed = sum(map(ord, str(cls.core_op))) + seed = sum(map(ord, str(cls.core_op) + cls.signature)) cls.rng = np.random.default_rng(seed) cls.params_sig, cls.outputs_sig = _parse_gufunc_signature(cls.signature) if cls.batcheable_axes is None: @@ -579,7 +579,10 @@ def test_solve(self, solve_fn, batched_A, batched_b): else: x = solve_fn(A, b, b_ndim=1) - mode = get_default_mode().excluding("batched_vector_b_solve_to_matrix_b_solve") + mode = get_default_mode().excluding( + "batched_vector_b_solve_to_matrix_b_solve", + "reuse_lu_decomposition_multiple_solves", + ) fn = function([In(A, mutable=True), In(b, mutable=True)], x, mode=mode) op = fn.maker.fgraph.outputs[0].owner.op