pytensor 2.37.0__tar.gz → 2.38.0__tar.gz

{pytensor-2.37.0/pytensor.egg-info → pytensor-2.38.0}/PKG-INFO

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

pytensor-2.38.0/doc/library/xtensor/signal.md

@@ -0,0 +1,7 @@
+(libdoc_xtensor_linalg)=
+# `xtensor.signal` -- Signal processing operations
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor.signal
+   :members:
+```

{pytensor-2.37.0 → pytensor-2.38.0}/pyproject.toml

@@ -163,6 +163,7 @@ lines-after-imports = 2
 "tests/link/numba/**/test_*.py" = ["E402"]
 "tests/link/pytorch/**/test_*.py" = ["E402"]
 "tests/link/mlx/**/test_*.py" = ["E402"]
+"tests/xtensor/**/*.py" = ["E402"]
 
 
 

{pytensor-2.37.0 → pytensor-2.38.0}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-01-23T09:04:26+0100",
+ "date": "2026-02-19T15:15:27+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "9834e96cee86d9f18436be5fa9e31c692bf6326f",
- "version": "2.37.0"
+ "full-revisionid": "87470065d89190405c9291235daf621d17de17ed",
+ "version": "2.38.0"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.37.0 → pytensor-2.38.0}/pytensor/graph/replace.py

@@ -208,19 +208,13 @@ def graph_replace(
 
 
 @singledispatch
-def _vectorize_node(op: Op, node: Apply, *batched_inputs) -> Apply:
+def _vectorize_node(op: Op, node: Apply, *batched_inputs) -> Apply | Sequence[Variable]:
     # Default implementation is provided in pytensor.tensor.blockwise
     raise NotImplementedError
 
 
-def vectorize_node(node: Apply, *batched_inputs) -> Apply:
-    """Returns vectorized version of node with new batched inputs."""
-    op = node.op
-    return _vectorize_node(op, node, *batched_inputs)
-
-
 def _vectorize_not_needed(op, node, *batched_inputs):
-    return op.make_node(*batched_inputs)
+    return op.make_node(*batched_inputs).outputs
 
 
 @overload
@@ -289,19 +283,40 @@ def vectorize_graph(
         # [array([-10., -11.]), array([10., 11.])]
 
     """
+    # TODO: Move this to tensor.vectorize, and make this helper type agnostic.
+    #
+    # This helper may dispatch to tensor.vectorize_graph or xtensor.vectorize_graph depending on the replacement types
+    # The behavior is distinct, because tensor vectorization depends on axis-position while xtensor depends on dimension labels
+    #
+    # xtensor.vectorize_graph will be able to handle batched inner tensor operations, while tensor.vectorize_graph won't,
+    # as it is by design unaware of xtensors and their semantics.
     if isinstance(outputs, Sequence):
         seq_outputs = outputs
     else:
         seq_outputs = [outputs]
 
+    if not all(
+        isinstance(key, Variable) and isinstance(value, Variable)
+        for key, value in replace.items()
+    ):
+        raise ValueError(f"Some of the replaced items are not Variables: {replace}")
+
     inputs = truncated_graph_inputs(seq_outputs, ancestors_to_include=replace.keys())
     new_inputs = [replace.get(inp, inp) for inp in inputs]
 
     vect_vars = dict(zip(inputs, new_inputs, strict=True))
     for node in toposort(seq_outputs, blockers=inputs):
         vect_inputs = [vect_vars.get(inp, inp) for inp in node.inputs]
-        vect_node = vectorize_node(node, *vect_inputs)
-        for output, vect_output in zip(node.outputs, vect_node.outputs, strict=True):
+
+        vect_node_or_outputs = _vectorize_node(node.op, node, *vect_inputs)
+        # Compatibility with the old API
+        vect_outputs = (
+            vect_node_or_outputs.outputs
+            if isinstance(vect_node_or_outputs, Apply)
+            else vect_node_or_outputs
+        )
+
+        for output, vect_output in zip(node.outputs, vect_outputs, strict=True):
             if output in vect_vars:
                 # This can happen when some outputs of a multi-output node are given a replacement,
                 # while some of the remaining outputs are still needed in the graph.

{pytensor-2.37.0 → pytensor-2.38.0}/pytensor/link/numba/dispatch/linalg/_LAPACK.py

@@ -946,3 +946,388 @@ class _LAPACK:
             fn(TRANA, TRANB, ISGN, M, N, A, LDA, B, LDB, C, LDC, SCALE, INFO)
 
         return trsyl
+
+    @classmethod
+    def numba_xgges(cls, dtype):
+        """
+        Compute generalized eigenvalues and, optionally, the left and/or right generalized Schur vectors of a pair
+        of real nonsymmetric matrices (A,B).
+
+        Called by scipy.linalg.qz and scipy.linalg.ordqz.
+        """
+        kind = get_blas_kind(dtype)
+        float_pointer = _get_nb_float_from_dtype(kind)
+        unique_func_name = f"scipy.lapack.{kind}gges"
+
+        @numba_basic.numba_njit
+        def get_gges_pointer():
+            with numba.objmode(ptr=types.intp):
+                ptr = get_lapack_ptr(dtype, "gges")
+            return ptr
+
+        if isinstance(dtype, Complex):
+            real_pointer = nb_f64p if dtype is nb_c128 else nb_f32p
+            gges_function_type = types.FunctionType(
+                types.void(
+                    nb_i32p,  # JOBVSL
+                    nb_i32p,  # JOBVSR
+                    nb_i32p,  # SORT
+                    nb_i32p,  # SELECT
+                    nb_i32p,  # N
+                    float_pointer,  # A
+                    nb_i32p,  # LDA
+                    float_pointer,  # B
+                    nb_i32p,  # LDB
+                    nb_i32p,  # SDIM
+                    float_pointer,  # ALPHA
+                    float_pointer,  # BETA
+                    float_pointer,  # VSL
+                    nb_i32p,  # LDVSL
+                    float_pointer,  # VSR
+                    nb_i32p,  # LDVSR
+                    float_pointer,  # WORK
+                    nb_i32p,  # LWORK
+                    real_pointer,  # RWORK
+                    nb_i32p,  # BWORK
+                    nb_i32p,  # INFO
+                )
+            )
+
+            @numba_basic.numba_njit
+            def gges(
+                JOBVSL,
+                JOBVSR,
+                SORT,
+                SELECT,
+                N,
+                A,
+                LDA,
+                B,
+                LDB,
+                SDIM,
+                ALPHA,
+                BETA,
+                VSL,
+                LDVSL,
+                VSR,
+                LDVSR,
+                WORK,
+                LWORK,
+                RWORK,
+                BWORK,
+                INFO,
+            ):
+                fn = _call_cached_ptr(
+                    get_ptr_func=get_gges_pointer,
+                    func_type_ref=gges_function_type,
+                    unique_func_name_lit=unique_func_name,
+                )
+                fn(
+                    JOBVSL,
+                    JOBVSR,
+                    SORT,
+                    SELECT,
+                    N,
+                    A,
+                    LDA,
+                    B,
+                    LDB,
+                    SDIM,
+                    ALPHA,
+                    BETA,
+                    VSL,
+                    LDVSL,
+                    VSR,
+                    LDVSR,
+                    WORK,
+                    LWORK,
+                    RWORK,
+                    BWORK,
+                    INFO,
+                )
+        else:  # Real case
+            gges_function_type = types.FunctionType(
+                types.void(
+                    nb_i32p,  # JOBVSL
+                    nb_i32p,  # JOBVSR
+                    nb_i32p,  # SORT
+                    nb_i32p,  # SELECT
+                    nb_i32p,  # N
+                    float_pointer,  # A
+                    nb_i32p,  # LDA
+                    float_pointer,  # B
+                    nb_i32p,  # LDB
+                    nb_i32p,  # SDIM
+                    float_pointer,  # ALPHAR
+                    float_pointer,  # ALPHAI
+                    float_pointer,  # BETA
+                    float_pointer,  # VSL
+                    nb_i32p,  # LDVSL
+                    float_pointer,  # VSR
+                    nb_i32p,  # LDVSR
+                    float_pointer,  # WORK
+                    nb_i32p,  # LWORK
+                    nb_i32p,  # BWORK
+                    nb_i32p,  # INFO
+                )
+            )
+
+            @numba_basic.numba_njit
+            def gges(
+                JOBVSL,
+                JOBVSR,
+                SORT,
+                SELECT,
+                N,
+                A,
+                LDA,
+                B,
+                LDB,
+                SDIM,
+                ALPHAR,
+                ALPHAI,
+                BETA,
+                VSL,
+                LDVSL,
+                VSR,
+                LDVSR,
+                WORK,
+                LWORK,
+                BWORK,
+                INFO,
+            ):
+                fn = _call_cached_ptr(
+                    get_ptr_func=get_gges_pointer,
+                    func_type_ref=gges_function_type,
+                    unique_func_name_lit=unique_func_name,
+                )
+                fn(
+                    JOBVSL,
+                    JOBVSR,
+                    SORT,
+                    SELECT,
+                    N,
+                    A,
+                    LDA,
+                    B,
+                    LDB,
+                    SDIM,
+                    ALPHAR,
+                    ALPHAI,
+                    BETA,
+                    VSL,
+                    LDVSL,
+                    VSR,
+                    LDVSR,
+                    WORK,
+                    LWORK,
+                    BWORK,
+                    INFO,
+                )
+
+        return gges
+
+    @classmethod
+    def numba_tgsen(cls, dtype):
+        """
+        Reorders the generalized Schur decomposition of a matrix pair (A, B) by their eigenvalues.
+
+        Output is sorted so that a selected cluster of eigenvalues appears in the leading diagonal blocks of the pair
+        (A,B). The leading columns of Q and Z form unitary bases of the corresponding left and right eigenspaces
+        (deflating subspaces). (A, B) must be in generalized Schur canonical form, that is, A and B are both upper
+        triangular.
+
+        Used by scipy.linalg.ordqz.
+        """
+        kind = get_blas_kind(dtype)
+        float_pointer = _get_nb_float_from_dtype(kind)
+        unique_func_name = f"scipy.lapack.{kind}tgsen"
+
+        @numba_basic.numba_njit
+        def get_tgsen_pointer():
+            with numba.objmode(ptr=types.intp):
+                ptr = get_lapack_ptr(dtype, "tgsen")
+            return ptr
+
+        if isinstance(dtype, Complex):
+            real_pointer = nb_f64p if dtype is nb_c128 else nb_f32p
+            tgsen_function_type = types.FunctionType(
+                types.void(
+                    nb_i32p,  # IJOB
+                    nb_i32p,  # WANTQ
+                    nb_i32p,  # WANTZ
+                    nb_i32p,  # SELECT
+                    nb_i32p,  # N
+                    float_pointer,  # A
+                    nb_i32p,  # LDA
+                    float_pointer,  # B
+                    nb_i32p,  # LDB
+                    float_pointer,  # alpha
+                    float_pointer,  # beta
+                    float_pointer,  # Q
+                    nb_i32p,  # LDQ
+                    float_pointer,  # Z
+                    nb_i32p,  # LDZ
+                    nb_i32p,  # M
+                    real_pointer,  # PL
+                    real_pointer,  # PR
+                    real_pointer,  # DIF
+                    float_pointer,  # WORK
+                    nb_i32p,  # LWORK
+                    nb_i32p,  # IWORK
+                    nb_i32p,  # LIWORK
+                    nb_i32p,  # INFO
+                )
+            )
+
+            @numba_basic.numba_njit
+            def tgsen(
+                IJOB,
+                WANTQ,
+                WANTZ,
+                SELECT,
+                N,
+                A,
+                LDA,
+                B,
+                LDB,
+                alpha,
+                beta,
+                Q,
+                LDQ,
+                Z,
+                LDZ,
+                M,
+                PL,
+                PR,
+                DIF,
+                WORK,
+                LWORK,
+                IWORK,
+                LIWORK,
+                INFO,
+            ):
+                fn = _call_cached_ptr(
+                    get_ptr_func=get_tgsen_pointer,
+                    func_type_ref=tgsen_function_type,
+                    unique_func_name_lit=unique_func_name,
+                )
+                fn(
+                    IJOB,
+                    WANTQ,
+                    WANTZ,
+                    SELECT,
+                    N,
+                    A,
+                    LDA,
+                    B,
+                    LDB,
+                    alpha,
+                    beta,
+                    Q,
+                    LDQ,
+                    Z,
+                    LDZ,
+                    M,
+                    PL,
+                    PR,
+                    DIF,
+                    WORK,
+                    LWORK,
+                    IWORK,
+                    LIWORK,
+                    INFO,
+                )
+        else:  # Real case
+            tgsen_function_type = types.FunctionType(
+                types.void(
+                    nb_i32p,  # IJOB
+                    nb_i32p,  # WANTQ
+                    nb_i32p,  # WANTZ
+                    nb_i32p,  # SELECT
+                    nb_i32p,  # N
+                    float_pointer,  # A
+                    nb_i32p,  # LDA
+                    float_pointer,  # B
+                    nb_i32p,  # LDB
+                    float_pointer,  # ALPHAR
+                    float_pointer,  # ALPHAI
+                    float_pointer,  # BETA
+                    float_pointer,  # Q
+                    nb_i32p,  # LDQ
+                    float_pointer,  # Z
+                    nb_i32p,  # LDZ
+                    nb_i32p,  # M
+                    float_pointer,  # PL
+                    float_pointer,  # PR
+                    float_pointer,  # DIF
+                    float_pointer,  # WORK
+                    nb_i32p,  # LWORK
+                    nb_i32p,  # IWORK
+                    nb_i32p,  # LIWORK
+                    nb_i32p,  # INFO
+                )
+            )
+
+            @numba_basic.numba_njit
+            def tgsen(
+                IJOB,
+                WANTQ,
+                WANTZ,
+                SELECT,
+                N,
+                A,
+                LDA,
+                B,
+                LDB,
+                ALPHAR,
+                ALPHAI,
+                BETA,
+                Q,
+                LDQ,
+                Z,
+                LDZ,
+                M,
+                PL,
+                PR,
+                DIF,
+                WORK,
+                LWORK,
+                IWORK,
+                LIWORK,
+                INFO,
+            ):
+                fn = _call_cached_ptr(
+                    get_ptr_func=get_tgsen_pointer,
+                    func_type_ref=tgsen_function_type,
+                    unique_func_name_lit=unique_func_name,
+                )
+                fn(
+                    IJOB,
+                    WANTQ,
+                    WANTZ,
+                    SELECT,
+                    N,
+                    A,
+                    LDA,
+                    B,
+                    LDB,
+                    ALPHAR,
+                    ALPHAI,
+                    BETA,
+                    Q,
+                    LDQ,
+                    Z,
+                    LDZ,
+                    M,
+                    PL,
+                    PR,
+                    DIF,
+                    WORK,
+                    LWORK,
+                    IWORK,
+                    LIWORK,
+                    INFO,
+                )
+
+        return tgsen