pytensor 2.31.0__tar.gz → 2.31.1__tar.gz

{pytensor-2.31.0/pytensor.egg-info → pytensor-2.31.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.31.0
+Version: 2.31.1
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause

{pytensor-2.31.0 → pytensor-2.31.1}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-05-19T16:51:53+0200",
+ "date": "2025-05-22T00:05:15+0800",
  "dirty": false,
  "error": null,
- "full-revisionid": "5ffe17ad7e1d2bc0841b69d0b57cfc9e142771c6",
- "version": "2.31.0"
+ "full-revisionid": "4829455bac15db5db652e3ab73e6baf08f6ea8b9",
+ "version": "2.31.1"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.31.0 → pytensor-2.31.1}/pytensor/compile/mode.py

@@ -490,6 +490,8 @@ PYTORCH = Mode(
             "fusion",
             "inplace",
             "scan_save_mem_prealloc",
+            "reuse_lu_decomposition_multiple_solves",
+            "scan_split_non_sequence_lu_decomposition_solve",
         ],
     ),
 )

{pytensor-2.31.0 → pytensor-2.31.1}/pytensor/scan/rewriting.py

@@ -2561,7 +2561,6 @@ scan_seqopt1.register(
     position=1,
 )
 
-
 scan_seqopt1.register(
     "scan_push_out_non_seq",
     in2out(scan_push_out_non_seq, ignore_newtrees=True),
@@ -2569,10 +2568,9 @@ scan_seqopt1.register(
     "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 @@ scan_seqopt1.register(
     "fast_run",
     "scan",
     "scan_pushout",
-    position=3,
+    position=4,
 )
 
 
@@ -2592,7 +2590,7 @@ scan_seqopt1.register(
     "more_mem",
     "scan",
     "scan_pushout",
-    position=4,
+    position=5,
 )
 
 
@@ -2605,7 +2603,7 @@ scan_seqopt1.register(
     "more_mem",
     "scan",
     "scan_pushout",
-    position=5,
+    position=6,
 )
 
 scan_eqopt2.register(

{pytensor-2.31.0 → pytensor-2.31.1}/pytensor/tensor/__init__.py

@@ -114,6 +114,7 @@ from pytensor.tensor import (
 
 
 # isort: off
+import pytensor.tensor._linalg
 from pytensor.tensor import linalg
 from pytensor.tensor import special
 from pytensor.tensor import signal

pytensor-2.31.1/pytensor/tensor/_linalg/__init__.py

@@ -0,0 +1,2 @@
+# Register rewrites
+import pytensor.tensor._linalg.solve

pytensor-2.31.1/pytensor/tensor/_linalg/solve/__init__.py

@@ -0,0 +1,2 @@
+# Register rewrites in the database
+import pytensor.tensor._linalg.solve.rewriting

pytensor-2.31.1/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,
+)

{pytensor-2.31.0 → pytensor-2.31.1}/pytensor/tensor/blas_c.py

@@ -344,6 +344,7 @@ def gemv_c_code(y, A, x, z, alpha, beta, fail, force_init_beta=False, params=Non
     """
     code = """
 
+    bool is_float;
     int elemsize;
     float fbeta;
     double dbeta;
@@ -361,11 +362,23 @@ def gemv_c_code(y, A, x, z, alpha, beta, fail, force_init_beta=False, params=Non
         %(fail)s;
     }
 
-    if  (PyArray_DESCR(%(y)s)->type_num == NPY_DOUBLE) { elemsize = 8; }
-    else if (PyArray_DESCR(%(y)s)->type_num == NPY_FLOAT) { elemsize = 4;}
+    if ((PyArray_DESCR(%(y)s)->type_num != PyArray_DESCR(%(x)s)->type_num)
+        || (PyArray_DESCR(%(y)s)->type_num != PyArray_DESCR(%(A)s)->type_num))
+    {
+        PyErr_SetString(PyExc_TypeError, "GEMV: dtypes of A, x, y do not match");
+        %(fail)s;
+    }
+    if  (PyArray_DESCR(%(y)s)->type_num == NPY_DOUBLE) {
+        is_float = 0;
+        elemsize = 8;
+    }
+    else if (PyArray_DESCR(%(y)s)->type_num == NPY_FLOAT) {
+        elemsize = 4;
+        is_float = 1;
+    }
     else {
-        PyErr_SetString(PyExc_NotImplementedError, "complex Gemv");
         %(fail)s;
+        PyErr_SetString(PyExc_NotImplementedError, "GEMV: Inputs must be float or double");
     }
 
     fbeta = dbeta = ((dtype_%(beta)s*)PyArray_DATA(%(beta)s))[0];
@@ -408,37 +421,40 @@ def gemv_c_code(y, A, x, z, alpha, beta, fail, force_init_beta=False, params=Non
             Py_INCREF(%(z)s);
         }
     }
+
     {
-        char TRANS = 'T';
-        char NOTRANS = 'N';
         int NA0 = PyArray_DIMS(%(A)s)[0];
         int NA1 = PyArray_DIMS(%(A)s)[1];
-        /* This formula is needed in the case where A is actually a row or
-         * column matrix, because BLAS sometimes insists that the strides:
-         *  - are not smaller than the number of elements in the array
-         *  - are not 0.
-         */
-        int SA0 = (NA0 > 1) ? (PyArray_STRIDES(%(A)s)[0] / elemsize) : (NA1 + 1);
-        int SA1 = (NA1 > 1) ? (PyArray_STRIDES(%(A)s)[1] / elemsize) : (NA0 + 1);
-        int Sz = PyArray_STRIDES(%(z)s)[0] / elemsize;
-        int Sx = PyArray_STRIDES(%(x)s)[0] / elemsize;
-
-        dtype_%(A)s* A_data = (dtype_%(A)s*) PyArray_DATA(%(A)s);
-        dtype_%(x)s* x_data = (dtype_%(x)s*) PyArray_DATA(%(x)s);
-        dtype_%(z)s* z_data = (dtype_%(z)s*) PyArray_DATA(%(z)s);
-        // gemv expects pointers to the beginning of memory arrays,
-        // but numpy provides a pointer to the first element,
-        // so when the stride is negative, we need to get the last one.
-        if (Sx < 0)
-            x_data += (NA1 - 1) * Sx;
-        if (Sz < 0)
-            z_data += (NA0 - 1) * Sz;
 
         if (NA0 * NA1)
         {
+            // Non-empty A matrix
+
+            /* In the case where A is actually a row or column matrix,
+             * the strides corresponding to the dummy dimension don't matter,
+             * but BLAS requires these to be no smaller than the number of elements in the array.
+             */
+            int SA0 = (NA0 > 1) ? (PyArray_STRIDES(%(A)s)[0] / elemsize) : NA1;
+            int SA1 = (NA1 > 1) ? (PyArray_STRIDES(%(A)s)[1] / elemsize) : NA0;
+            int Sz = PyArray_STRIDES(%(z)s)[0] / elemsize;
+            int Sx = PyArray_STRIDES(%(x)s)[0] / elemsize;
+
+            dtype_%(A)s* A_data = (dtype_%(A)s*) PyArray_DATA(%(A)s);
+            dtype_%(x)s* x_data = (dtype_%(x)s*) PyArray_DATA(%(x)s);
+            dtype_%(z)s* z_data = (dtype_%(z)s*) PyArray_DATA(%(z)s);
+
+            // gemv expects pointers to the beginning of memory arrays,
+            // but numpy provides a pointer to the first element,
+            // so when the stride is negative, we need to get the last one.
+            if (Sx < 0)
+                x_data += (NA1 - 1) * Sx;
+            if (Sz < 0)
+                z_data += (NA0 - 1) * Sz;
+
             if ( ((SA0 < 0) || (SA1 < 0)) && (abs(SA0) == 1 || (abs(SA1) == 1)) )
             {
                 // We can treat the array A as C-or F-contiguous by changing the order of iteration
+                // printf("GEMV: Iterating in reverse NA0=%%d, NA1=%%d, SA0=%%d, SA1=%%d\\n", NA0, NA1, SA0, SA1);
                 if (SA0 < 0){
                     A_data += (NA0 -1) * SA0;  // Jump to first row
                     SA0 = -SA0;  // Iterate over rows in reverse
@@ -452,27 +468,45 @@ def gemv_c_code(y, A, x, z, alpha, beta, fail, force_init_beta=False, params=Non
             } else if ((SA0 < 0) || (SA1 < 0) || ((SA0 != 1) && (SA1 != 1)))
             {
                 // Array isn't contiguous, we have to make a copy
-                // - if the copy is too long, maybe call vector/vector dot on
-                //   each row instead
-                // printf("GEMV: Making a copy SA0=%%d, SA1=%%d\\n", SA0, SA1);
+                // - if the copy is too long, maybe call vector/vector dot on each row instead
+                // printf("GEMV: Making a copy NA0=%%d, NA1=%%d, SA0=%%d, SA1=%%d\\n", NA0, NA1, SA0, SA1);
                 npy_intp dims[2];
                 dims[0] = NA0;
                 dims[1] = NA1;
-
-                PyArrayObject * A_copy = (PyArrayObject *) PyArray_Copy(
-                                                                   %(A)s);
+                PyArrayObject * A_copy = (PyArrayObject *) PyArray_Copy(%(A)s);
                 if (!A_copy)
                     %(fail)s
                 Py_XDECREF(%(A)s);
                 %(A)s = A_copy;
-                SA0 = (NA0 > 1) ? (PyArray_STRIDES(%(A)s)[0] / elemsize) : (NA1 + 1);
-                SA1 = (NA1 > 1) ? (PyArray_STRIDES(%(A)s)[1] / elemsize) : (NA0 + 1);
+                SA0 = (NA0 > 1) ? (PyArray_STRIDES(%(A)s)[0] / elemsize) : NA1;
+                SA1 = (NA1 > 1) ? (PyArray_STRIDES(%(A)s)[1] / elemsize) : NA0;
                 A_data = (dtype_%(A)s*) PyArray_DATA(%(A)s);
             }
+            //else {printf("GEMV: Using the original array NA0=%%d, NA1=%%d, SA0=%%d, SA1=%%d\\n", NA0, NA1, SA0, SA1);}
 
-            if (SA0 == 1)
+            if (NA0 == 1)
+            {
+                // Vector-vector dot product, it seems faster to avoid GEMV
+                dtype_%(alpha)s alpha = ((dtype_%(alpha)s*)PyArray_DATA(%(alpha)s))[0];
+
+                if (is_float)
+                {
+                    z_data[0] *= fbeta;
+                    z_data[0] += alpha * sdot_(&NA1,  (float*)(A_data), &SA1,
+                                              (float*)x_data, &Sx);
+                }
+                else
+                {
+                    z_data[0] *= dbeta;
+                    z_data[0] += alpha * ddot_(&NA1,  (double*)(A_data), &SA1,
+                                              (double*)x_data, &Sx);
+                }
+            }
+            else if (SA0 == 1)
             {
-                if (PyArray_DESCR(%(A)s)->type_num == NPY_FLOAT)
+                // F-contiguous
+                char NOTRANS = 'N';
+                if (is_float)
                 {
                     float alpha = ((dtype_%(alpha)s*)PyArray_DATA(%(alpha)s))[0];
                     sgemv_(&NOTRANS, &NA0, &NA1,
@@ -482,7 +516,7 @@ def gemv_c_code(y, A, x, z, alpha, beta, fail, force_init_beta=False, params=Non
                         &fbeta,
                         (float*)z_data, &Sz);
                 }
-                else if (PyArray_DESCR(%(A)s)->type_num == NPY_DOUBLE)
+                else
                 {
                     double alpha = ((dtype_%(alpha)s*)PyArray_DATA(%(alpha)s))[0];
                     dgemv_(&NOTRANS, &NA0, &NA1,
@@ -492,97 +526,39 @@ def gemv_c_code(y, A, x, z, alpha, beta, fail, force_init_beta=False, params=Non
                         &dbeta,
                         (double*)z_data, &Sz);
                 }
-                else
-                {
-                    PyErr_SetString(PyExc_AssertionError,
-                                    "neither float nor double dtype");
-                    %(fail)s
-                }
             }
             else if (SA1 == 1)
             {
-                if (PyArray_DESCR(%(A)s)->type_num == NPY_FLOAT)
+                // C-contiguous
+                char TRANS = 'T';
+                if (is_float)
                 {
                     float alpha = ((dtype_%(alpha)s*)PyArray_DATA(%(alpha)s))[0];
-
-                    // Check for vector-vector dot (NA0 == 1). The code may work
-                    // for SA1 != 1 as well, but has not been tested for this case,
-                    // so SA1 == 1 is required for safety.
-                    if (NA0 == 1 && SA1 == 1)
-                    {
-                        if (fbeta != 0.f) {
-                          z_data[0] = fbeta*z_data[0];
-                        } else {
-                          z_data[0] = 0.f;
-                        }
-                        z_data[0] += alpha*sdot_(&NA1,
-                              (float*)(A_data), &SA1,
-                              (float*)x_data, &Sx);
-                    }
-                    else
-                    {
-                        sgemv_(&TRANS, &NA1, &NA0,
-                            &alpha,
-                            (float*)(A_data), &SA0,
-                            (float*)x_data, &Sx,
-                            &fbeta,
-                            (float*)z_data, &Sz);
-                    }
-                }
-                else if (PyArray_DESCR(%(A)s)->type_num == NPY_DOUBLE)
-                {
-                    double alpha = ((dtype_%(alpha)s*)PyArray_DATA(%(alpha)s))[0];
-
-                    // Check for vector-vector dot (NA0 == 1). The code may work
-                    // for SA1 != 1 as well, but has not been tested for this case,
-                    // so SA1 == 1 is required for safety.
-                    if (NA0 == 1 && SA1 == 1)
-                    {
-                        if (dbeta != 0.) {
-                          z_data[0] = dbeta*z_data[0];
-                        } else {
-                          z_data[0] = 0.;
-                        }
-                        z_data[0] += alpha*ddot_(&NA1,
-                              (double*)(A_data), &SA1,
-                              (double*)x_data, &Sx);
-                    }
-                    else
-                    {
-                        dgemv_(&TRANS, &NA1, &NA0,
-                            &alpha,
-                            (double*)(A_data), &SA0,
-                            (double*)x_data, &Sx,
-                            &dbeta,
-                            (double*)z_data, &Sz);
-                    }
+                    sgemv_(&TRANS, &NA1, &NA0,
+                        &alpha,
+                        (float*)(A_data), &SA0,
+                        (float*)x_data, &Sx,
+                        &fbeta,
+                        (float*)z_data, &Sz);
                 }
                 else
                 {
-                    PyErr_SetString(PyExc_AssertionError,
-                                    "neither float nor double dtype");
-                    %(fail)s
+                    double alpha = ((dtype_%(alpha)s*)PyArray_DATA(%(alpha)s))[0];
+                    dgemv_(&TRANS, &NA1, &NA0,
+                        &alpha,
+                        (double*)(A_data), &SA0,
+                        (double*)x_data, &Sx,
+                        &dbeta,
+                        (double*)z_data, &Sz);
                 }
             }
             else
             {
                 PyErr_SetString(PyExc_AssertionError,
-                    "xx is a double-strided matrix, and should have been "
-                    "copied into a memory-contiguous one.");
+                                "A is neither C nor F-contiguous, it should have been copied into a memory-contiguous array;");
                 %(fail)s
             }
         }
-        else if (dbeta != 1.0)
-        {
-            // the matrix has at least one dim of length 0
-            // so we do this loop, which either iterates over 0 elements
-            // or else it does the right thing for length-0 A.
-            dtype_%(z)s * zptr = (dtype_%(z)s*)(PyArray_DATA(%(z)s));
-            for (int i = 0; i < NA0; ++i)
-            {
-                zptr[i * Sz] = (dbeta == 0.0 ? 0.0 : zptr[i * Sz] * dbeta);
-            }
-        }
     }
     """
     return code % locals()
@@ -613,7 +589,7 @@ class CGemv(BaseBLAS, Gemv):
         return code
 
     def c_code_cache_version(self):
-        return (15, blas_header_version(), check_force_gemv_init())
+        return (16, blas_header_version(), check_force_gemv_init())
 
 
 cgemv_inplace = CGemv(inplace=True)

{pytensor-2.31.0 → pytensor-2.31.1}/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
 

{pytensor-2.31.0 → pytensor-2.31.1}/pytensor/tensor/rewriting/subtensor.py

@@ -25,6 +25,7 @@ from pytensor.tensor.basic import (
     alloc,
     cast,
     concatenate,
+    expand_dims,
     get_scalar_constant_value,
     get_underlying_scalar_constant_value,
     register_infer_shape,
@@ -1576,7 +1577,15 @@ def local_blockwise_advanced_inc_subtensor(fgraph, node):
         x = alloc(x, *batch_shape, *core_shape)
 
     new_idxs = [slice(None)] * batch_ndim + new_idxs
-    symbolic_idxs = x[tuple(new_idxs)].owner.inputs[1:]
+    x_view = x[tuple(new_idxs)]
+
+    # We need to introduce any implicit expand_dims on core dimension of y
+    y_core_ndim = y.type.ndim - batch_ndim
+    if (missing_y_core_ndim := x_view.type.ndim - batch_ndim - y_core_ndim) > 0:
+        missing_axes = tuple(range(batch_ndim, batch_ndim + missing_y_core_ndim))
+        y = expand_dims(y, missing_axes)
+
+    symbolic_idxs = x_view.owner.inputs[1:]
     new_out = op.core_op.make_node(x, y, *symbolic_idxs).outputs
     copy_stack_trace(node.outputs, new_out)
     return new_out

{pytensor-2.31.0 → pytensor-2.31.1/pytensor.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.31.0
+Version: 2.31.1
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause

{pytensor-2.31.0 → pytensor-2.31.1}/pytensor.egg-info/SOURCES.txt

@@ -394,6 +394,9 @@ pytensor/tensor/utils.py
 pytensor/tensor/var.py
 pytensor/tensor/variable.py
 pytensor/tensor/xlogx.py
+pytensor/tensor/_linalg/__init__.py
+pytensor/tensor/_linalg/solve/__init__.py
+pytensor/tensor/_linalg/solve/rewriting.py
 pytensor/tensor/c_code/alt_blas_common.h
 pytensor/tensor/c_code/alt_blas_template.c
 pytensor/tensor/c_code/dimshuffle.c