pytensor 2.36.3__tar.gz → 2.38.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.36.3
+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.36.3 → pytensor-2.38.0}/doc/extending/creating_an_op.rst

@@ -495,7 +495,7 @@ If both outputs are disconnected PyTensor will not bother calling the :meth:`L_o
 
     from pytensor.graph.op import Op
     from pytensor.graph.basic import Apply
-    from pytensor.gradient import DisconnectedType
+    from pytensor.gradient import DisconnectedType, disconnected_type
 
     class TransposeAndSumOp(Op):
         __props__ = ()
@@ -539,13 +539,13 @@ If both outputs are disconnected PyTensor will not bother calling the :meth:`L_o
             out1_grad, out2_grad = output_grads
 
             if isinstance(out1_grad.type, DisconnectedType):
-                x_grad = DisconnectedType()()
+                x_grad = disconnected_type()
             else:
                 # Transpose the last two dimensions of the output gradient
                 x_grad = pt.swapaxes(out1_grad, -1, -2)
 
             if isinstance(out2_grad.type, DisconnectedType):
-                y_grad = DisconnectedType()()
+                y_grad = disconnected_type()
             else:
                 # Broadcast the output gradient to the same shape as y
                 y_grad = pt.broadcast_to(pt.expand_dims(out2_grad, -1), y.shape)

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.36.3 → pytensor-2.38.0}/pyproject.toml

@@ -149,27 +149,22 @@ lines-after-imports = 2
 # TODO: Get rid of these:
 "**/__init__.py" = ["F401", "E402", "F403"]
 "pytensor/tensor/linalg.py" = ["F403"]
+# Modules that use print-statements, skip "T201"
 "pytensor/link/c/cmodule.py" = ["PTH", "T201"]
 "pytensor/misc/elemwise_time_test.py" = ["T201"]
 "pytensor/misc/elemwise_openmp_speedup.py" = ["T201"]
 "pytensor/misc/check_duplicate_key.py" = ["T201"]
 "pytensor/misc/check_blas.py" = ["T201"]
 "pytensor/bin/pytensor_cache.py" = ["T201"]
-# For the tests we skip `E402` because `pytest.importorskip` is used:
-"tests/link/jax/test_scalar.py" = ["E402"]
-"tests/link/jax/test_tensor_basic.py" = ["E402"]
-"tests/link/numba/test_basic.py" = ["E402"]
-"tests/link/numba/test_cython_support.py" = ["E402"]
-"tests/link/numba/test_performance.py" = ["E402"]
-"tests/link/numba/test_sparse.py" = ["E402"]
-"tests/link/numba/test_tensor_basic.py" = ["E402"]
-"tests/tensor/test_math_scipy.py" = ["E402"]
-"tests/sparse/test_basic.py" = ["E402"]
-"tests/sparse/test_sp2.py" = ["E402"]
-"tests/sparse/test_utils.py" = ["E402"]
-"tests/sparse/sandbox/test_sp.py" = ["E402", "F401"]
 "tests/compile/test_monitormode.py" = ["T201"]
 "scripts/run_mypy.py" = ["T201"]
+# Test modules of optional backends that use `pytest.importorskip`, skip "E402"
+"tests/link/jax/**/test_*.py" = ["E402"]
+"tests/link/numba/**/test_*.py" = ["E402"]
+"tests/link/pytorch/**/test_*.py" = ["E402"]
+"tests/link/mlx/**/test_*.py" = ["E402"]
+"tests/xtensor/**/*.py" = ["E402"]
+
 
 
 [tool.mypy]

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

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-01-07T19:04:45+0100",
+ "date": "2026-02-19T15:15:27+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "ac11da625fc3155790bc0b4794807232b73de0b8",
- "version": "2.36.3"
+ "full-revisionid": "87470065d89190405c9291235daf621d17de17ed",
+ "version": "2.38.0"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.36.3 → pytensor-2.38.0}/pytensor/breakpoint.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from pytensor.gradient import DisconnectedType
+from pytensor.gradient import disconnected_type
 from pytensor.graph.basic import Apply, Variable
 from pytensor.graph.op import Op
 from pytensor.tensor.basic import as_tensor_variable
@@ -142,7 +142,7 @@ class PdbBreakpoint(Op):
                 output_storage[i][0] = inputs[i + 1]
 
     def grad(self, inputs, output_gradients):
-        return [DisconnectedType()(), *output_gradients]
+        return [disconnected_type(), *output_gradients]
 
     def infer_shape(self, fgraph, inputs, input_shapes):
         # Return the shape of every input but the condition (first input)

{pytensor-2.36.3 → pytensor-2.38.0}/pytensor/gradient.py

@@ -494,22 +494,25 @@ def Lop(
         coordinates of the tensor elements.
         If `f` is a list/tuple, then return a list/tuple with the results.
     """
-    if not isinstance(eval_points, list | tuple):
-        _eval_points: list[Variable] = [pytensor.tensor.as_tensor_variable(eval_points)]
-    else:
-        _eval_points = [pytensor.tensor.as_tensor_variable(x) for x in eval_points]
+    from pytensor.tensor import as_tensor_variable
 
-    if not isinstance(f, list | tuple):
-        _f: list[Variable] = [pytensor.tensor.as_tensor_variable(f)]
-    else:
-        _f = [pytensor.tensor.as_tensor_variable(x) for x in f]
+    if not isinstance(eval_points, Sequence):
+        eval_points = [eval_points]
+    _eval_points = [
+        x if isinstance(x, Variable) else as_tensor_variable(x) for x in eval_points
+    ]
+
+    if not isinstance(f, Sequence):
+        f = [f]
+    _f = [x if isinstance(x, Variable) else as_tensor_variable(x) for x in f]
 
     grads = list(_eval_points)
 
-    if not isinstance(wrt, list | tuple):
-        _wrt: list[Variable] = [pytensor.tensor.as_tensor_variable(wrt)]
-    else:
-        _wrt = [pytensor.tensor.as_tensor_variable(x) for x in wrt]
+    using_list = isinstance(wrt, list)
+    using_tuple = isinstance(wrt, tuple)
+    if not isinstance(wrt, Sequence):
+        wrt = [wrt]
+    _wrt = [x if isinstance(x, Variable) else as_tensor_variable(x) for x in wrt]
 
     assert len(_f) == len(grads)
     known = dict(zip(_f, grads, strict=True))
@@ -523,8 +526,6 @@ def Lop(
         return_disconnected=return_disconnected,
     )
 
-    using_list = isinstance(wrt, list)
-    using_tuple = isinstance(wrt, tuple)
     return as_list_or_tuple(using_list, using_tuple, ret)
 
 

{pytensor-2.36.3 → 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.36.3 → pytensor-2.38.0}/pytensor/link/basic.py

@@ -284,7 +284,7 @@ class PerformLinker(LocalLinker):
     """
 
     required_rewrites: tuple[str, ...] = ("minimum_compile", "py_only")
-    incompatible_rewrites: tuple[str, ...] = ("cxx",)
+    incompatible_rewrites: tuple[str, ...] = ("cxx_only",)
 
     def __init__(
         self, allow_gc: bool | None = None, schedule: Callable | None = None

{pytensor-2.36.3 → pytensor-2.38.0}/pytensor/link/jax/dispatch/slinalg.py

@@ -3,6 +3,7 @@ import warnings
 import jax
 
 from pytensor.link.jax.dispatch.basic import jax_funcify
+from pytensor.tensor._linalg.solve.linear_control import SolveSylvester
 from pytensor.tensor.slinalg import (
     LU,
     QR,
@@ -13,6 +14,7 @@ from pytensor.tensor.slinalg import (
     Expm,
     LUFactor,
     PivotToPermutations,
+    Schur,
     Solve,
     SolveTriangular,
 )
@@ -92,7 +94,6 @@ def jax_funcify_Solve(op, **kwargs):
 def jax_funcify_SolveTriangular(op, **kwargs):
     lower = op.lower
     unit_diagonal = op.unit_diagonal
-    check_finite = op.check_finite
 
     def solve_triangular(A, b):
         return jax.scipy.linalg.solve_triangular(
@@ -101,7 +102,7 @@ def jax_funcify_SolveTriangular(op, **kwargs):
             lower=lower,
             trans=0,  # this is handled by explicitly transposing A, so it will always be 0 when we get to here.
             unit_diagonal=unit_diagonal,
-            check_finite=check_finite,
+            check_finite=False,
         )
 
     return solve_triangular
@@ -132,27 +133,23 @@ def jax_funcify_PivotToPermutation(op, **kwargs):
 def jax_funcify_LU(op, **kwargs):
     permute_l = op.permute_l
     p_indices = op.p_indices
-    check_finite = op.check_finite
 
     if p_indices:
         raise ValueError("JAX does not support the p_indices argument")
 
     def lu(*inputs):
-        return jax.scipy.linalg.lu(
-            *inputs, permute_l=permute_l, check_finite=check_finite
-        )
+        return jax.scipy.linalg.lu(*inputs, permute_l=permute_l, check_finite=False)
 
     return lu
 
 
 @jax_funcify.register(LUFactor)
 def jax_funcify_LUFactor(op, **kwargs):
-    check_finite = op.check_finite
     overwrite_a = op.overwrite_a
 
     def lu_factor(a):
         return jax.scipy.linalg.lu_factor(
-            a, check_finite=check_finite, overwrite_a=overwrite_a
+            a, check_finite=False, overwrite_a=overwrite_a
         )
 
     return lu_factor
@@ -161,12 +158,11 @@ def jax_funcify_LUFactor(op, **kwargs):
 @jax_funcify.register(CholeskySolve)
 def jax_funcify_ChoSolve(op, **kwargs):
     lower = op.lower
-    check_finite = op.check_finite
     overwrite_b = op.overwrite_b
 
     def cho_solve(c, b):
         return jax.scipy.linalg.cho_solve(
-            (c, lower), b, check_finite=check_finite, overwrite_b=overwrite_b
+            (c, lower), b, check_finite=False, overwrite_b=overwrite_b
         )
 
     return cho_solve
@@ -189,3 +185,27 @@ def jax_funcify_Expm(op, **kwargs):
         return jax.scipy.linalg.expm(x)
 
     return expm
+
+
+@jax_funcify.register(Schur)
+def jax_funcify_Schur(op, **kwargs):
+    output = op.output
+
+    if op.sort is not None:
+        warnings.warn(
+            "jax.scipy.linalg.schur only supports sort=None. The sort argument is ignored."
+        )
+
+    def schur(a):
+        T, Z = jax.scipy.linalg.schur(a, output=output)
+        return T, Z
+
+    return schur
+
+
+@jax_funcify.register(SolveSylvester)
+def jax_funcify_SolveSylsterer(op, **kwargs):
+    def solve_sylvester(a, b, c):
+        return jax.scipy.linalg.solve_sylvester(a, b, c)
+
+    return solve_sylvester

{pytensor-2.36.3 → pytensor-2.38.0}/pytensor/link/jax/linker.py

@@ -14,7 +14,7 @@ class JAXLinker(JITLinker):
         "jax",
     )  # TODO: Distinguish between optional "jax" and "minimum_compile_jax"
     incompatible_rewrites = (
-        "cxx",
+        "cxx_only",
         "BlasOpt",
         "local_careduce_fusion",
         "scan_save_mem_prealloc",

{pytensor-2.36.3 → pytensor-2.38.0}/pytensor/link/numba/dispatch/basic.py

@@ -14,14 +14,13 @@ from pytensor.graph.basic import Apply, Constant, Variable
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.type import Type
 from pytensor.link.numba.cache import compile_numba_function_src, hash_from_pickle_dump
-from pytensor.link.numba.dispatch.sparse import CSCMatrixType, CSRMatrixType
 from pytensor.link.utils import (
     fgraph_to_python,
 )
 from pytensor.scalar.basic import ScalarType
 from pytensor.sparse import SparseTensorType
 from pytensor.tensor.random.type import RandomGeneratorType
-from pytensor.tensor.type import TensorType
+from pytensor.tensor.type import DenseTensorType
 from pytensor.tensor.utils import hash_from_ndarray
 from pytensor.typed_list import TypedListType
 
@@ -112,7 +111,7 @@ def get_numba_type(
         Return Numba scalars for zero dimensional :class:`TensorType`\s.
     """
 
-    if isinstance(pytensor_type, TensorType):
+    if isinstance(pytensor_type, DenseTensorType):
         dtype = pytensor_type.numpy_dtype
         numba_dtype = numba.from_dtype(dtype)
         if force_scalar or (
@@ -125,18 +124,26 @@ def get_numba_type(
         numba_dtype = numba.from_dtype(dtype)
         return numba_dtype
     elif isinstance(pytensor_type, SparseTensorType):
-        dtype = pytensor_type.numpy_dtype
-        numba_dtype = numba.from_dtype(dtype)
+        from pytensor.link.numba.dispatch.sparse.variable import (
+            CSCMatrixType,
+            CSRMatrixType,
+        )
+
+        data_array = numba.types.Array(
+            numba.from_dtype(pytensor_type.numpy_dtype), 1, layout
+        )
+        indices_array = numba.types.Array(numba.from_dtype(np.int32), 1, layout)
+        indptr_array = numba.types.Array(numba.from_dtype(np.int32), 1, layout)
         if pytensor_type.format == "csr":
-            return CSRMatrixType(numba_dtype)
+            return CSRMatrixType(data_array, indices_array, indptr_array)
         if pytensor_type.format == "csc":
-            return CSCMatrixType(numba_dtype)
+            return CSCMatrixType(data_array, indices_array, indptr_array)
     elif isinstance(pytensor_type, RandomGeneratorType):
         return numba.types.NumPyRandomGeneratorType("NumPyRandomGeneratorType")
     elif isinstance(pytensor_type, TypedListType):
         return numba.types.List(get_numba_type(pytensor_type.ttype))
-    else:
-        raise NotImplementedError(f"Numba type not implemented for {pytensor_type}")
+
+    raise NotImplementedError(f"Numba type not implemented for {pytensor_type}")
 
 
 def create_numba_signature(

{pytensor-2.36.3 → pytensor-2.38.0}/pytensor/link/numba/dispatch/blockwise.py

@@ -86,6 +86,7 @@ def numba_funcify_Blockwise(op: BlockwiseWithCoreShape, node, **kwargs):
                 output_bc_patterns,
                 output_dtypes,
                 inplace_pattern,
+                False,  # allow_core_scalar
                 (),  # constant_inputs
                 inputs,
                 tuple_core_shapes,
@@ -98,6 +99,7 @@ def numba_funcify_Blockwise(op: BlockwiseWithCoreShape, node, **kwargs):
         # If the core op cannot be cached, the Blockwise wrapper cannot be cached either
         blockwise_key = None
     else:
+        blockwise_cache_version = 1
         blockwise_key = "_".join(
             map(
                 str,
@@ -108,6 +110,7 @@ def numba_funcify_Blockwise(op: BlockwiseWithCoreShape, node, **kwargs):
                     blockwise_op.signature,
                     input_bc_patterns,
                     core_op_key,
+                    blockwise_cache_version,
                 ),
             )
         )

{pytensor-2.36.3 → pytensor-2.38.0}/pytensor/link/numba/dispatch/compile_ops.py

@@ -68,7 +68,7 @@ def numba_funcify_OpFromGraph(op, node=None, **kwargs):
     output_specs = [Out(o, borrow=False) for o in fgraph.outputs]
     insert_deepcopy(fgraph, wrapped_inputs=input_specs, wrapped_outputs=output_specs)
     fgraph_fn, fgraph_cache_key = numba_funcify_and_cache_key(
-        fgraph, squeeze_output=True, **kwargs
+        fgraph, squeeze_output=True, fgraph_name="numba_ofg", **kwargs
     )
 
     if fgraph_cache_key is None:

{pytensor-2.36.3 → pytensor-2.38.0}/pytensor/link/numba/dispatch/elemwise.py

@@ -365,6 +365,7 @@ def numba_funcify_Elemwise(op, node, **kwargs):
                 output_bc_patterns_enc,
                 output_dtypes_enc,
                 inplace_pattern_enc,
+                True,  # allow_core_scalar
                 (),  # constant_inputs
                 inputs,
                 core_output_shapes,  # core_shapes