pytensor 2.31.1__tar.gz → 2.31.2__tar.gz

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

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

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-05-22T00:05:15+0800",
+ "date": "2025-05-24T11:34:23+0800",
  "dirty": false,
  "error": null,
- "full-revisionid": "4829455bac15db5db652e3ab73e6baf08f6ea8b9",
- "version": "2.31.1"
+ "full-revisionid": "2d414d4129882861b867aaf0d9abcf3135e3d5eb",
+ "version": "2.31.2"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.31.1 → 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.1 → 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.1 → 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.1 → 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.1 → 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.1 → 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.1 → pytensor-2.31.2}/pytensor/tensor/_linalg/solve/rewriting.py

@@ -14,16 +14,22 @@ from pytensor.tensor.slinalg import Solve, lu_factor, lu_solve
 from pytensor.tensor.variable import TensorVariable
 
 
-def decompose_A(A, assume_a):
+def decompose_A(A, assume_a, check_finite):
     if assume_a == "gen":
-        return lu_factor(A, check_finite=False)
+        return lu_factor(A, check_finite=check_finite)
     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)
+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
 
@@ -102,14 +108,19 @@ def _split_lu_solve_steps(
         ):
             return None
 
-    A_decomp = decompose_A(A, assume_a=assume_a)
+    # 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
-        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
+            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)

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

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

@@ -336,11 +336,12 @@ cger_no_inplace = CGer(False)
 # ##### ####### #######
 
 
-def gemv_c_code(y, A, x, z, alpha, beta, fail, force_init_beta=False, params=None):
+def gemv_c_code(y, A, x, z, alpha, beta, fail, must_initialize_y=False, params=None):
     """
     z <- beta * y + alpha * dot(A, x)
 
     where A is a matrix, y and x are vectors (ergo z is vector)
+    z = y if inplace else y.copy()
     """
     code = """
 
@@ -400,17 +401,11 @@ def gemv_c_code(y, A, x, z, alpha, beta, fail, force_init_beta=False, params=Non
         }
         if (dbeta != 0)
         {
+            // If dbeta is zero, we avoid doing the copy
             if (PyArray_CopyInto(%(z)s, %(y)s) != 0) {
                 %(fail)s
             }
         }
-        else if (%(force_init_beta)d)
-        {
-            PyObject *zero = PyFloat_FromDouble(0.);
-            if (zero == NULL) %(fail)s;
-            if (PyArray_FillWithScalar(%(z)s, zero) != 0) %(fail)s;
-            Py_DECREF(zero);
-        }
     }
     else
     {
@@ -422,6 +417,20 @@ def gemv_c_code(y, A, x, z, alpha, beta, fail, force_init_beta=False, params=Non
         }
     }
 
+    if (%(must_initialize_y)d && dbeta == 0)
+    {
+        // 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.
+        PyObject *zero = PyFloat_FromDouble(0.);
+        if (zero == NULL) %(fail)s;
+        if (PyArray_FillWithScalar(%(z)s, zero) != 0) %(fail)s;
+        Py_DECREF(zero);
+    }
+
     {
         int NA0 = PyArray_DIMS(%(A)s)[0];
         int NA1 = PyArray_DIMS(%(A)s)[1];
@@ -491,13 +500,13 @@ def gemv_c_code(y, A, x, z, alpha, beta, fail, force_init_beta=False, params=Non
 
                 if (is_float)
                 {
-                    z_data[0] *= fbeta;
+                    z_data[0] = dbeta != 0 ? dbeta * z_data[0] : 0.f;
                     z_data[0] += alpha * sdot_(&NA1,  (float*)(A_data), &SA1,
                                               (float*)x_data, &Sx);
                 }
                 else
                 {
-                    z_data[0] *= dbeta;
+                    z_data[0] = dbeta != 0 ? dbeta * z_data[0] : 0.;
                     z_data[0] += alpha * ddot_(&NA1,  (double*)(A_data), &SA1,
                                               (double*)x_data, &Sx);
                 }
@@ -583,21 +592,21 @@ class CGemv(BaseBLAS, Gemv):
             alpha,
             beta,
             fail=sub["fail"],
-            force_init_beta=check_force_gemv_init(),
+            must_initialize_y=must_initialize_y_gemv(),
             params=sub["params"],
         )
         return code
 
     def c_code_cache_version(self):
-        return (16, blas_header_version(), check_force_gemv_init())
+        return (17, blas_header_version(), must_initialize_y_gemv())
 
 
 cgemv_inplace = CGemv(inplace=True)
 cgemv_no_inplace = CGemv(inplace=False)
 
 
-def check_force_gemv_init():
-    if check_force_gemv_init._force_init_beta is None:
+def must_initialize_y_gemv():
+    if must_initialize_y_gemv._force_init_beta is None:
         from pytensor.link.c.cmodule import GCC_compiler
 
         """
@@ -643,13 +652,13 @@ int main() {
         )
         if res:
             if res[0]:
-                check_force_gemv_init._force_init_beta = res[1]
+                must_initialize_y_gemv._force_init_beta = res[1]
             else:
-                check_force_gemv_init._force_init_beta = False
+                must_initialize_y_gemv._force_init_beta = False
         else:
-            check_force_gemv_init._force_init_beta = False
+            must_initialize_y_gemv._force_init_beta = False
 
-    return check_force_gemv_init._force_init_beta
+    return must_initialize_y_gemv._force_init_beta
 
 
-check_force_gemv_init._force_init_beta = None
+must_initialize_y_gemv._force_init_beta = None

{pytensor-2.31.1 → pytensor-2.31.2}/pytensor/tensor/blockwise.py

@@ -16,6 +16,7 @@ from pytensor.graph.replace import (
     _vectorize_not_needed,
     vectorize_graph,
 )
+from pytensor.link.c.op import COp
 from pytensor.scalar import ScalarType
 from pytensor.tensor import as_tensor_variable
 from pytensor.tensor.shape import shape_padleft
@@ -43,7 +44,18 @@ def _vectorize_node_perform(
     """
 
     storage_map = {var: [None] for var in core_node.inputs + core_node.outputs}
-    core_thunk = core_node.op.make_thunk(core_node, storage_map, None, [], impl=impl)
+    try:
+        core_thunk = core_node.op.make_thunk(
+            core_node, storage_map, None, [], impl=impl
+        )
+    except NotImplementedError:
+        if impl == "c":
+            # Try again with py impl
+            core_thunk = core_node.op.make_thunk(
+                core_node, storage_map, None, [], impl="py"
+            )
+        else:
+            raise
     single_in = len(core_node.inputs) == 1
     core_input_storage = [storage_map[inp] for inp in core_node.inputs]
     core_output_storage = [storage_map[out] for out in core_node.outputs]
@@ -128,7 +140,7 @@ def _check_runtime_broadcast_core(numerical_inputs, batch_bcast_patterns, batch_
             )
 
 
-class Blockwise(Op):
+class Blockwise(COp):
     """Generalizes a core `Op` to work with batched dimensions.
 
     TODO: Dispatch JAX (should be easy with the vectorize macro)
@@ -483,6 +495,14 @@ class Blockwise(Op):
         else:
             return self.name
 
+    def c_code(self, *args, **kwargs):
+        # Blockwise is a C_Op just so we can propagate compilation mode to the inner Op.
+        # It doesn't itself have a C implementation yet.
+        raise NotImplementedError()
+
+    def c_code_cache_version(self):
+        return (-1,)
+
 
 @_vectorize_node.register(Op)
 def vectorize_node_fallback(op: Op, node: Apply, *bached_inputs) -> Apply:

{pytensor-2.31.1 → pytensor-2.31.2}/pytensor/tensor/rewriting/blas.py

@@ -700,7 +700,7 @@ def local_dot22_to_ger_or_gemv(fgraph, node):
             new_out = [rval]
         elif xb[0] and yb[1]:
             # x and y are both vectors so this qualifies for a sdot / ddot
-            # TODO: PyTensor doesn't have a sdot, but gemv is better than _dot22
+            # PyTensor's CGemv will call sdot/ddot at runtime, the Scipy Gemv may not
             xv = x.dimshuffle(1)
             zeros = ptb.AllocEmpty(x.dtype)(1)
             rval = gemv_no_inplace(zeros, one, y.T, xv, zero)

{pytensor-2.31.1 → pytensor-2.31.2}/pytensor/tensor/signal/conv.py

@@ -2,7 +2,8 @@ from typing import TYPE_CHECKING, Literal, cast
 
 from numpy import convolve as numpy_convolve
 
-from pytensor.graph import Apply, Op
+from pytensor.graph import Apply
+from pytensor.link.c.op import COp
 from pytensor.scalar.basic import upcast
 from pytensor.tensor.basic import as_tensor_variable, join, zeros
 from pytensor.tensor.blockwise import Blockwise
@@ -15,7 +16,7 @@ if TYPE_CHECKING:
     from pytensor.tensor import TensorLike
 
 
-class Convolve1d(Op):
+class Convolve1d(COp):
     __props__ = ("mode",)
     gufunc_signature = "(n),(k)->(o)"
 
@@ -86,6 +87,87 @@ class Convolve1d(Op):
 
         return [in1_bar, in2_bar]
 
+    def c_code_cache_version(self):
+        return (1,)
+
+    def c_code(self, node, name, inputs, outputs, sub):
+        # raise NotImplementedError()
+        in1, in2 = inputs
+        [out] = outputs
+        mode_str = self.mode
+
+        if mode_str == "full":
+            np_mode_val = 2  # NPY_CONVOLVE_FULL
+        elif mode_str == "valid":
+            np_mode_val = 0  # NPY_CONVOLVE_VALID
+        else:
+            # This case should ideally be prevented by __init__ or make_node
+            raise ValueError(f"Unsupported mode {mode_str}")
+
+        code = f"""
+        {{
+            PyArrayObject* in2_flipped_view = NULL;
+
+            if (PyArray_NDIM({in1}) != 1 || PyArray_NDIM({in2}) != 1) {{
+                PyErr_SetString(PyExc_ValueError, "Convolve1d C code expects 1D arrays.");
+                {sub['fail']};
+            }}
+
+            npy_intp n_in2 = PyArray_DIM({in2}, 0);
+
+            // Create a reversed view of in2
+            if (n_in2 == 0) {{
+                PyErr_SetString(PyExc_ValueError, "Convolve1d: second input (kernel) cannot be empty.");
+                {sub['fail']};
+            }} else {{
+                npy_intp view_dims[1];
+                view_dims[0] = n_in2;
+
+                npy_intp view_strides[1];
+                view_strides[0] = -PyArray_STRIDES({in2})[0];
+
+                void* view_data = (char*)PyArray_DATA({in2}) + (n_in2 - 1) * PyArray_STRIDES({in2})[0];
+
+                Py_INCREF(PyArray_DESCR({in2}));
+                in2_flipped_view = (PyArrayObject*)PyArray_NewFromDescr(
+                    Py_TYPE({in2}),
+                    PyArray_DESCR({in2}),
+                    1,  // ndim
+                    view_dims,
+                    view_strides,
+                    view_data,
+                    (PyArray_FLAGS({in2}) & ~NPY_ARRAY_WRITEABLE),
+                    NULL
+                );
+
+                if (!in2_flipped_view) {{
+                    PyErr_SetString(PyExc_RuntimeError, "Failed to create flipped kernel view for Convolve1d.");
+                    {sub['fail']};
+                }}
+
+                Py_INCREF({in2});
+                if (PyArray_SetBaseObject(in2_flipped_view, (PyObject*){in2}) < 0) {{
+                    Py_DECREF({in2}); // SetBaseObject failed, release the extra INCREF
+                    Py_DECREF(in2_flipped_view);
+                    in2_flipped_view = NULL;
+                    PyErr_SetString(PyExc_RuntimeError, "Failed to set base object for flipped kernel view in Convolve1d.");
+                    {sub['fail']};
+                }}
+                PyArray_UpdateFlags(in2_flipped_view, (NPY_ARRAY_C_CONTIGUOUS | NPY_ARRAY_F_CONTIGUOUS));
+            }}
+
+            // TODO: Use lower level implementation that allows reusing the output buffer
+            Py_XDECREF({out});
+            {out} = (PyArrayObject*) PyArray_Correlate2((PyObject*){in1}, (PyObject*)in2_flipped_view, {np_mode_val});
+            Py_XDECREF(in2_flipped_view); // Clean up the view if correlate fails
+            if (!{out}) {{
+                // PyArray_Correlate already set an error
+                {sub['fail']};
+            }}
+        }}
+        """
+        return code
+
 
 def convolve1d(
     in1: "TensorLike",

{pytensor-2.31.1 → pytensor-2.31.2}/pytensor/tensor/type.py

@@ -793,7 +793,7 @@ def tensor(
         try:
             # Help catching errors with the new tensor API
             # Many single letter strings are valid sctypes
-            if str(name) == "floatX" or (len(str(name)) > 1 and np.dtype(name).type):
+            if str(name) == "floatX" or (len(str(name)) > 2 and np.dtype(name).type):
                 raise ValueError(
                     f"The first and only positional argument of tensor is now `name`. Got {name}.\n"
                     "This name looks like a dtype, which you should pass as a keyword argument only."

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.31.1
+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