pytensor 2.36.3__tar.gz → 2.37.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.36.3
+Version: 2.37.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.37.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.36.3 → pytensor-2.37.0}/pyproject.toml

@@ -149,27 +149,21 @@ 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"]
+
 
 
 [tool.mypy]

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

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-01-07T19:04:45+0100",
+ "date": "2026-01-23T09:04:26+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "ac11da625fc3155790bc0b4794807232b73de0b8",
- "version": "2.36.3"
+ "full-revisionid": "9834e96cee86d9f18436be5fa9e31c692bf6326f",
+ "version": "2.37.0"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.36.3 → pytensor-2.37.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.37.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.37.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.37.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.37.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.37.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.37.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.37.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.37.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