pytensor 2.31.7__tar.gz → 2.32.0__tar.gz

{pytensor-2.31.7/pytensor.egg-info → pytensor-2.32.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.31.7
+Version: 2.32.0
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause
@@ -55,8 +55,6 @@ Provides-Extra: rtd
 Requires-Dist: sphinx<6,>=5.1.0; extra == "rtd"
 Requires-Dist: pygments; extra == "rtd"
 Requires-Dist: pydot; extra == "rtd"
-Requires-Dist: pydot2; extra == "rtd"
-Requires-Dist: pydot-ng; extra == "rtd"
 Provides-Extra: jax
 Requires-Dist: jax; extra == "jax"
 Requires-Dist: jaxlib; extra == "jax"

{pytensor-2.31.7 → pytensor-2.32.0}/doc/library/xtensor/type.md

@@ -14,7 +14,6 @@
 
 ```{eval-rst}
 .. automodule:: pytensor.xtensor.type
-   :noindex:
    :members: XTensorType, XTensorVariable, XTensorConstant
 ```
 

{pytensor-2.31.7 → pytensor-2.32.0}/pyproject.toml

@@ -78,7 +78,7 @@ tests = [
     "pytest-mock",
     "pytest-sphinx",
 ]
-rtd = ["sphinx>=5.1.0,<6", "pygments", "pydot", "pydot2", "pydot-ng"]
+rtd = ["sphinx>=5.1.0,<6", "pygments", "pydot"]
 jax = ["jax", "jaxlib"]
 numba = ["numba>=0.57", "llvmlite"]
 

{pytensor-2.31.7 → pytensor-2.32.0}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-07-08T17:42:30+0200",
+ "date": "2025-07-26T12:45:37+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "4ce092fe4df857a2bf88a6a367b4a1592e10c0ee",
- "version": "2.31.7"
+ "full-revisionid": "6f8bb55521559b0f5ebb14ef536becca7c70f2a7",
+ "version": "2.32.0"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.31.7 → pytensor-2.32.0}/pytensor/graph/rewriting/basic.py

@@ -1550,6 +1550,7 @@ class PatternNodeRewriter(NodeRewriter):
         tracks=(),
         get_nodes=None,
         values_eq_approx=None,
+        allow_cast=True,
     ):
         """
 
@@ -1572,6 +1573,10 @@ class PatternNodeRewriter(NodeRewriter):
             If you provide `tracks`, you must provide this parameter. It must be a
             function that takes the tracked node and returns a list of nodes on
             which we will try this rewrite.
+        values_eq_approx
+            TODO
+        allow_cast
+            Automatically cast the output of the rewrite whenever new and old types differ
 
         Notes
         -----
@@ -1586,6 +1591,7 @@ class PatternNodeRewriter(NodeRewriter):
         self.in_pattern = convert_strs_to_vars(in_pattern, var_map=var_map)
         self.out_pattern = convert_strs_to_vars(out_pattern, var_map=var_map)
         self.values_eq_approx = values_eq_approx
+        self.allow_cast = allow_cast
         if isinstance(in_pattern, list | tuple):
             self.op = self.in_pattern[0]
         elif isinstance(in_pattern, dict):
@@ -1630,6 +1636,10 @@ class PatternNodeRewriter(NodeRewriter):
         if node.op != self.op:
             return False
 
+        if len(node.outputs) != 1:
+            # PatternNodeRewriter doesn't support replacing multi-output nodes
+            return False
+
         s = unify(self.in_pattern, node.out)
 
         if s is False:
@@ -1652,19 +1662,20 @@ class PatternNodeRewriter(NodeRewriter):
             ):
                 return False
 
-        if ret.owner:
+        [old_out] = node.outputs
+        if not old_out.type.is_super(ret.type):
+            # Type doesn't match
             if not (
-                len(node.outputs) == len(ret.owner.outputs)
-                and all(
-                    o.type.is_super(new_o.type)
-                    for o, new_o in zip(node.outputs, ret.owner.outputs, strict=True)
-                )
+                self.allow_cast
+                and isinstance(old_out.type, pytensor.tensor.TensorType)
+                and isinstance(ret.type, pytensor.tensor.TensorType)
             ):
                 return False
-        else:
-            # ret is just an input variable
-            assert len(node.outputs) == 1
-            if not node.outputs[0].type.is_super(ret.type):
+
+            # Try to cast tensors
+            ret = ret.astype(old_out.type.dtype)
+            if not old_out.type.is_super(ret.type):
+                # Still doesn't match
                 return False
 
         return [ret]

{pytensor-2.31.7 → pytensor-2.32.0}/pytensor/link/jax/dispatch/basic.py

@@ -10,10 +10,11 @@ from pytensor.compile import JAX
 from pytensor.compile.builders import OpFromGraph
 from pytensor.compile.ops import DeepCopyOp, TypeCastingOp
 from pytensor.configdefaults import config
+from pytensor.graph import Constant
 from pytensor.graph.fg import FunctionGraph
 from pytensor.ifelse import IfElse
 from pytensor.link.utils import fgraph_to_python
-from pytensor.raise_op import Assert, CheckAndRaise
+from pytensor.raise_op import CheckAndRaise
 
 
 if config.floatX == "float64":
@@ -73,11 +74,14 @@ def jax_funcify_IfElse(op, **kwargs):
     return ifelse
 
 
-@jax_funcify.register(Assert)
 @jax_funcify.register(CheckAndRaise)
-def jax_funcify_CheckAndRaise(op, **kwargs):
+def jax_funcify_CheckAndRaise(op, node, **kwargs):
+    conds = node.inputs[1:]
+    if any(isinstance(cond, Constant) and not bool(cond.data) for cond in conds):
+        raise op.exc_type(op.msg)
+
     warnings.warn(
-        f"""Skipping `CheckAndRaise` Op (assertion: {op.msg}) as JAX tracing would remove it.""",
+        f"""Skipping {op} Op (assertion: {op.msg}) as JAX tracing would remove it.""",
         stacklevel=2,
     )
 

{pytensor-2.31.7 → pytensor-2.32.0}/pytensor/link/jax/dispatch/nlinalg.py

@@ -9,7 +9,6 @@ from pytensor.tensor.nlinalg import (
     KroneckerProduct,
     MatrixInverse,
     MatrixPinv,
-    QRFull,
     SLogDet,
 )
 
@@ -67,16 +66,6 @@ def jax_funcify_MatrixInverse(op, **kwargs):
     return matrix_inverse
 
 
-@jax_funcify.register(QRFull)
-def jax_funcify_QRFull(op, **kwargs):
-    mode = op.mode
-
-    def qr_full(x, mode=mode):
-        return jnp.linalg.qr(x, mode=mode)
-
-    return qr_full
-
-
 @jax_funcify.register(MatrixPinv)
 def jax_funcify_Pinv(op, **kwargs):
     def pinv(x):

{pytensor-2.31.7 → pytensor-2.32.0}/pytensor/link/jax/dispatch/slinalg.py

@@ -5,6 +5,7 @@ import jax
 from pytensor.link.jax.dispatch.basic import jax_funcify
 from pytensor.tensor.slinalg import (
     LU,
+    QR,
     BlockDiagonal,
     Cholesky,
     CholeskySolve,
@@ -168,3 +169,13 @@ def jax_funcify_ChoSolve(op, **kwargs):
         )
 
     return cho_solve
+
+
+@jax_funcify.register(QR)
+def jax_funcify_QR(op, **kwargs):
+    mode = op.mode
+
+    def qr(x, mode=mode):
+        return jax.scipy.linalg.qr(x, mode=mode)
+
+    return qr

{pytensor-2.31.7 → pytensor-2.32.0}/pytensor/link/numba/dispatch/basic.py

@@ -402,24 +402,22 @@ def numba_funcify_DeepCopyOp(op, node, **kwargs):
     return deepcopyop
 
 
-@numba_njit
-def makeslice(*x):
-    return slice(*x)
-
-
 @numba_funcify.register(MakeSlice)
 def numba_funcify_MakeSlice(op, **kwargs):
-    return global_numba_func(makeslice)
-
+    @numba_njit
+    def makeslice(*x):
+        return slice(*x)
 
-@numba_njit
-def shape(x):
-    return np.asarray(np.shape(x))
+    return makeslice
 
 
 @numba_funcify.register(Shape)
 def numba_funcify_Shape(op, **kwargs):
-    return global_numba_func(shape)
+    @numba_njit
+    def shape(x):
+        return np.asarray(np.shape(x))
+
+    return shape
 
 
 @numba_funcify.register(Shape_i)

{pytensor-2.31.7 → pytensor-2.32.0}/pytensor/link/numba/dispatch/linalg/_LAPACK.py

@@ -283,7 +283,6 @@ class _LAPACK:
 
         Called by scipy.linalg.lu_solve
         """
-        ...
         lapack_ptr, float_pointer = _get_lapack_ptr_and_ptr_type(dtype, "getrs")
         functype = ctypes.CFUNCTYPE(
             None,
@@ -457,3 +456,90 @@ class _LAPACK:
             _ptr_int,  # INFO
         )
         return functype(lapack_ptr)
+
+    @classmethod
+    def numba_xgeqrf(cls, dtype):
+        """
+        Compute the QR factorization of a general M-by-N matrix A.
+
+        Used in QR decomposition (no pivoting).
+        """
+        lapack_ptr, float_pointer = _get_lapack_ptr_and_ptr_type(dtype, "geqrf")
+        functype = ctypes.CFUNCTYPE(
+            None,
+            _ptr_int,  # M
+            _ptr_int,  # N
+            float_pointer,  # A
+            _ptr_int,  # LDA
+            float_pointer,  # TAU
+            float_pointer,  # WORK
+            _ptr_int,  # LWORK
+            _ptr_int,  # INFO
+        )
+        return functype(lapack_ptr)
+
+    @classmethod
+    def numba_xgeqp3(cls, dtype):
+        """
+        Compute the QR factorization with column pivoting of a general M-by-N matrix A.
+
+        Used in QR decomposition with pivoting.
+        """
+        lapack_ptr, float_pointer = _get_lapack_ptr_and_ptr_type(dtype, "geqp3")
+        functype = ctypes.CFUNCTYPE(
+            None,
+            _ptr_int,  # M
+            _ptr_int,  # N
+            float_pointer,  # A
+            _ptr_int,  # LDA
+            _ptr_int,  # JPVT
+            float_pointer,  # TAU
+            float_pointer,  # WORK
+            _ptr_int,  # LWORK
+            _ptr_int,  # INFO
+        )
+        return functype(lapack_ptr)
+
+    @classmethod
+    def numba_xorgqr(cls, dtype):
+        """
+        Generate the orthogonal matrix Q from a QR factorization (real types).
+
+        Used in QR decomposition to form Q.
+        """
+        lapack_ptr, float_pointer = _get_lapack_ptr_and_ptr_type(dtype, "orgqr")
+        functype = ctypes.CFUNCTYPE(
+            None,
+            _ptr_int,  # M
+            _ptr_int,  # N
+            _ptr_int,  # K
+            float_pointer,  # A
+            _ptr_int,  # LDA
+            float_pointer,  # TAU
+            float_pointer,  # WORK
+            _ptr_int,  # LWORK
+            _ptr_int,  # INFO
+        )
+        return functype(lapack_ptr)
+
+    @classmethod
+    def numba_xungqr(cls, dtype):
+        """
+        Generate the unitary matrix Q from a QR factorization (complex types).
+
+        Used in QR decomposition to form Q for complex types.
+        """
+        lapack_ptr, float_pointer = _get_lapack_ptr_and_ptr_type(dtype, "ungqr")
+        functype = ctypes.CFUNCTYPE(
+            None,
+            _ptr_int,  # M
+            _ptr_int,  # N
+            _ptr_int,  # K
+            float_pointer,  # A
+            _ptr_int,  # LDA
+            float_pointer,  # TAU
+            float_pointer,  # WORK
+            _ptr_int,  # LWORK
+            _ptr_int,  # INFO
+        )
+        return functype(lapack_ptr)