pytensor 2.31.0__tar.gz → 2.31.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.31.0
+Version: 2.31.2
 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.2}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-05-19T16:51:53+0200",
+ "date": "2025-05-24T11:34:23+0800",
  "dirty": false,
  "error": null,
- "full-revisionid": "5ffe17ad7e1d2bc0841b69d0b57cfc9e142771c6",
- "version": "2.31.0"
+ "full-revisionid": "2d414d4129882861b867aaf0d9abcf3135e3d5eb",
+ "version": "2.31.2"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.31.0 → pytensor-2.31.2}/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.2}/pytensor/link/jax/dispatch/slinalg.py

@@ -46,6 +46,7 @@ def jax_funcify_Cholesky(op, **kwargs):
 def jax_funcify_Solve(op, **kwargs):
     assume_a = op.assume_a
     lower = op.lower
+    b_is_vec = op.b_ndim == 1
 
     if assume_a == "tridiagonal":
         # jax.scipy.solve does not yet support tridiagonal matrices
@@ -54,7 +55,20 @@ def jax_funcify_Solve(op, **kwargs):
             dl = jax.numpy.diagonal(a, offset=-1, axis1=-2, axis2=-1)
             d = jax.numpy.diagonal(a, offset=0, axis1=-2, axis2=-1)
             du = jax.numpy.diagonal(a, offset=1, axis1=-2, axis2=-1)
-            return jax.lax.linalg.tridiagonal_solve(dl, d, du, b, lower=lower)
+
+            # jax requires dl and du to have the same shape as d
+            dl = jax.numpy.pad(dl, (1, 0))
+            du = jax.numpy.pad(du, (0, 1))
+
+            if b_is_vec:
+                b = jax.numpy.expand_dims(b, -1)
+
+            res = jax.lax.linalg.tridiagonal_solve(dl, d, du, b)
+
+            if b_is_vec:
+                return jax.numpy.squeeze(res, -1)
+
+            return res
 
     else:
         if assume_a not in ("gen", "sym", "her", "pos"):

{pytensor-2.31.0 → pytensor-2.31.2}/pytensor/link/numba/dispatch/basic.py

@@ -31,7 +31,6 @@ from pytensor.link.utils import (
     fgraph_to_python,
 )
 from pytensor.scalar.basic import ScalarType
-from pytensor.scalar.math import Softplus
 from pytensor.sparse import SparseTensorType
 from pytensor.tensor.basic import Nonzero
 from pytensor.tensor.blas import BatchedDot
@@ -466,7 +465,7 @@ def numba_funcify_ArgSortOp(op, node, **kwargs):
             axis = axis.item()
 
             Y = np.swapaxes(X, axis, 0)
-            result = np.empty_like(Y)
+            result = np.empty_like(Y, dtype="int64")
 
             indices = list(np.ndindex(Y.shape[1:]))
 
@@ -607,25 +606,6 @@ def numba_funcify_Dot(op, node, **kwargs):
     return dot
 
 
-@numba_funcify.register(Softplus)
-def numba_funcify_Softplus(op, node, **kwargs):
-    x_dtype = np.dtype(node.inputs[0].dtype)
-
-    @numba_njit
-    def softplus(x):
-        if x < -37.0:
-            value = np.exp(x)
-        elif x < 18.0:
-            value = np.log1p(np.exp(x))
-        elif x < 33.3:
-            value = x + np.exp(-x)
-        else:
-            value = x
-        return direct_cast(value, x_dtype)
-
-    return softplus
-
-
 @numba_funcify.register(Solve)
 def numba_funcify_Solve(op, node, **kwargs):
     assume_a = op.assume_a
@@ -689,11 +669,6 @@ def numba_funcify_BatchedDot(op, node, **kwargs):
     return batched_dot
 
 
-# NOTE: The remaining `pytensor.tensor.blas` `Op`s appear unnecessary, because
-# they're only used to optimize basic `Dot` nodes, and those GEMV and GEMM
-# optimizations are apparently already performed by Numba
-
-
 @numba_funcify.register(IfElse)
 def numba_funcify_IfElse(op, **kwargs):
     n_outs = op.n_outs

{pytensor-2.31.0 → pytensor-2.31.2}/pytensor/link/numba/dispatch/elemwise.py

@@ -561,7 +561,7 @@ def numba_funcify_Argmax(op, node, **kwargs):
 
         @numba_basic.numba_njit(inline="always")
         def argmax(x):
-            return 0
+            return np.array(0, dtype="int64")
 
     else:
         axes = tuple(int(ax) for ax in axis)

{pytensor-2.31.0 → pytensor-2.31.2}/pytensor/link/numba/dispatch/linalg/decomposition/lu.py

@@ -30,7 +30,7 @@ def _lu_factor_to_lu(a, dtype, overwrite_a):
     # Fortran is 1 indexed, so we need to subtract 1 from the IPIV array
     IPIV = IPIV - 1
     p_inv = _pivot_to_permutation(IPIV, dtype=dtype)
-    perm = np.argsort(p_inv)
+    perm = np.argsort(p_inv).astype("int32")
 
     return perm, L, U
 

{pytensor-2.31.0 → pytensor-2.31.2}/pytensor/link/numba/dispatch/nlinalg.py

@@ -52,7 +52,7 @@ def numba_funcify_Det(op, node, **kwargs):
 
     @numba_basic.numba_njit(inline="always")
     def det(x):
-        return numba_basic.direct_cast(np.linalg.det(inputs_cast(x)), out_dtype)
+        return np.array(np.linalg.det(inputs_cast(x))).astype(out_dtype)
 
     return det
 
@@ -68,8 +68,8 @@ def numba_funcify_SLogDet(op, node, **kwargs):
     def slogdet(x):
         sign, det = np.linalg.slogdet(inputs_cast(x))
         return (
-            numba_basic.direct_cast(sign, out_dtype_1),
-            numba_basic.direct_cast(det, out_dtype_2),
+            np.array(sign).astype(out_dtype_1),
+            np.array(det).astype(out_dtype_2),
         )
 
     return slogdet

{pytensor-2.31.0 → pytensor-2.31.2}/pytensor/link/numba/dispatch/scalar.py

@@ -28,7 +28,7 @@ from pytensor.scalar.basic import (
     Second,
     Switch,
 )
-from pytensor.scalar.math import Erf, Erfc, GammaLn, Log1mexp, Sigmoid
+from pytensor.scalar.math import Erf, Erfc, GammaLn, Log1mexp, Sigmoid, Softplus
 
 
 @numba_funcify.register(ScalarOp)
@@ -312,3 +312,22 @@ def erfc(x):
 @numba_funcify.register(Erfc)
 def numba_funcify_Erfc(op, **kwargs):
     return numba_basic.global_numba_func(erfc)
+
+
+@numba_funcify.register(Softplus)
+def numba_funcify_Softplus(op, node, **kwargs):
+    out_dtype = np.dtype(node.outputs[0].type.dtype)
+
+    @numba_basic.numba_njit
+    def softplus(x):
+        if x < -37.0:
+            value = np.exp(x)
+        elif x < 18.0:
+            value = np.log1p(np.exp(x))
+        elif x < 33.3:
+            value = x + np.exp(-x)
+        else:
+            value = x
+        return numba_basic.direct_cast(value, out_dtype)
+
+    return softplus

{pytensor-2.31.0 → pytensor-2.31.2}/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.2}/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.2/pytensor/tensor/_linalg/__init__.py

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

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

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

pytensor-2.31.2/pytensor/tensor/_linalg/solve/rewriting.py

@@ -0,0 +1,209 @@
+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, check_finite):
+    if assume_a == "gen":
+        return lu_factor(A, check_finite=check_finite)
+    else:
+        raise NotImplementedError
+
+
+def solve_lu_decomposed_system(A_decomp, b, transposed=False, *, core_solve_op: Solve):
+    if core_solve_op.assume_a == "gen":
+        return lu_solve(
+            A_decomp,
+            b,
+            trans=transposed,
+            b_ndim=core_solve_op.b_ndim,
+            check_finite=core_solve_op.check_finite,
+        )
+    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
+
+    # If any Op had check_finite=True, we also do it for the LU decomposition
+    check_finite_decomp = False
+    for client, _ in A_solve_clients_and_transpose:
+        if client.op.core_op.check_finite:
+            check_finite_decomp = True
+            break
+    A_decomp = decompose_A(A, assume_a=assume_a, check_finite=check_finite_decomp)
+
+    replacements = {}
+    for client, transposed in A_solve_clients_and_transpose:
+        _, b = client.inputs
+        new_x = solve_lu_decomposed_system(
+            A_decomp, b, transposed=transposed, core_solve_op=client.op.core_op
+        )
+        [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.2}/pytensor/tensor/blas.py

@@ -113,23 +113,22 @@ from pytensor.tensor.type import DenseTensorType, tensor
 _logger = logging.getLogger("pytensor.tensor.blas")
 
 
-# If check_init_y() == True we need to initialize y when beta == 0.
-def check_init_y():
-    # TODO: What is going on here?
+def must_initialize_y_gemv():
+    # Check whether Scipy GEMV could output nan if y in not initialized
     from scipy.linalg.blas import get_blas_funcs
 
-    if check_init_y._result is None:
-        y = float("NaN") * np.ones((2,))
+    if must_initialize_y_gemv._result is None:
+        y = np.full((2,), np.nan)
         x = np.ones((2,))
         A = np.ones((2, 2))
         gemv = get_blas_funcs("gemv", dtype=y.dtype)
         gemv(1.0, A.T, x, 0.0, y, overwrite_y=True, trans=True)
-        check_init_y._result = np.isnan(y).any()
+        must_initialize_y_gemv._result = np.isnan(y).any()
 
-    return check_init_y._result
+    return must_initialize_y_gemv._result
 
 
-check_init_y._result = None  # type: ignore
+must_initialize_y_gemv._result = None  # type: ignore
 
 
 class Gemv(Op):
@@ -197,7 +196,13 @@ class Gemv(Op):
                     f"(beta * y + alpha * dot(A, x)). y: {y.shape}, A: {A.shape}, x: {x.shape}"
                 )
 
-            if beta == 0 and check_init_y():
+            if beta == 0 and must_initialize_y_gemv():
+                # Most BLAS implementations of GEMV ignore y=nan when beta=0
+                # PyTensor considers that the correct behavior,
+                # and even exploits it to avoid copying or initializing outputs.
+                # By deciding to exploit this, however, it becomes our responsibility
+                # to ensure the behavior even in the rare cases BLAS deviates,
+                # or users will get errors, even for graphs that had no nan to begin with.
                 y.fill(0)
 
             # Here I suppose that A is in c order. If we don't make it