pytensor 2.31.1__tar.gz → 2.31.3__tar.gz

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

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

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-05-22T00:05:15+0800",
+ "date": "2025-06-02T15:09:01+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "4829455bac15db5db652e3ab73e6baf08f6ea8b9",
- "version": "2.31.1"
+ "full-revisionid": "ff98ab8f8d2e6b78539a9225bde5e18a49ae6b81",
+ "version": "2.31.3"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.31.1 → pytensor-2.31.3}/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.3}/pytensor/link/jax/dispatch/tensor_basic.py

@@ -87,14 +87,7 @@ def jax_funcify_Join(op, **kwargs):
     def join(axis, *tensors):
         # tensors could also be tuples, and in this case they don't have a ndim
         tensors = [jnp.asarray(tensor) for tensor in tensors]
-        view = op.view
-        if (view != -1) and all(
-            tensor.shape[axis] == 0 for tensor in tensors[0:view] + tensors[view + 1 :]
-        ):
-            return tensors[view]
-
-        else:
-            return jnp.concatenate(tensors, axis=axis)
+        return jnp.concatenate(tensors, axis=axis)
 
     return join
 

{pytensor-2.31.1 → pytensor-2.31.3}/pytensor/link/jax/linker.py

@@ -9,8 +9,13 @@ from pytensor.link.basic import JITLinker
 class JAXLinker(JITLinker):
     """A `Linker` that JIT-compiles NumPy-based operations using JAX."""
 
+    def __init__(self, *args, **kwargs):
+        self.scalar_shape_inputs: tuple[int] = ()  # type: ignore[annotation-unchecked]
+        super().__init__(*args, **kwargs)
+
     def fgraph_convert(self, fgraph, input_storage, storage_map, **kwargs):
         from pytensor.link.jax.dispatch import jax_funcify
+        from pytensor.link.jax.dispatch.shape import JAXShapeTuple
         from pytensor.tensor.random.type import RandomType
 
         shared_rng_inputs = [
@@ -64,6 +69,23 @@ class JAXLinker(JITLinker):
                 fgraph.inputs.remove(new_inp)
                 fgraph.inputs.insert(old_inp_fgrap_index, new_inp)
 
+        fgraph_inputs = fgraph.inputs
+        clients = fgraph.clients
+        # Detect scalar shape inputs that are used only in JAXShapeTuple nodes
+        scalar_shape_inputs = [
+            inp
+            for node in fgraph.apply_nodes
+            if isinstance(node.op, JAXShapeTuple)
+            for inp in node.inputs
+            if inp in fgraph_inputs
+            and all(
+                isinstance(cl_node.op, JAXShapeTuple) for cl_node, _ in clients[inp]
+            )
+        ]
+        self.scalar_shape_inputs = tuple(
+            fgraph_inputs.index(inp) for inp in scalar_shape_inputs
+        )
+
         return jax_funcify(
             fgraph, input_storage=input_storage, storage_map=storage_map, **kwargs
         )
@@ -71,7 +93,22 @@ class JAXLinker(JITLinker):
     def jit_compile(self, fn):
         import jax
 
-        return jax.jit(fn)
+        jit_fn = jax.jit(fn, static_argnums=self.scalar_shape_inputs)
+
+        if not self.scalar_shape_inputs:
+            return jit_fn
+
+        def convert_scalar_shape_inputs(
+            *args, scalar_shape_inputs=set(self.scalar_shape_inputs)
+        ):
+            return jit_fn(
+                *(
+                    int(arg) if i in scalar_shape_inputs else arg
+                    for i, arg in enumerate(args)
+                )
+            )
+
+        return convert_scalar_shape_inputs
 
     def create_thunk_inputs(self, storage_map):
         from pytensor.link.jax.dispatch import jax_typify

{pytensor-2.31.1 → pytensor-2.31.3}/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:]))
 
@@ -566,18 +565,19 @@ def numba_funcify_SpecifyShape(op, node, **kwargs):
 def int_to_float_fn(inputs, out_dtype):
     """Create a Numba function that converts integer and boolean ``ndarray``s to floats."""
 
-    if all(
-        input.type.numpy_dtype == np.dtype(out_dtype) for input in inputs
-    ) and isinstance(np.dtype(out_dtype), np.floating):
+    if (
+        all(inp.type.dtype == out_dtype for inp in inputs)
+        and np.dtype(out_dtype).kind == "f"
+    ):
 
-        @numba_njit
+        @numba_njit(inline="always")
         def inputs_cast(x):
             return x
 
-    elif any(i.type.numpy_dtype.kind in "ib" for i in inputs):
+    elif any(i.type.numpy_dtype.kind in "uib" for i in inputs):
         args_dtype = np.dtype(f"f{out_dtype.itemsize}")
 
-        @numba_njit
+        @numba_njit(inline="always")
         def inputs_cast(x):
             return x.astype(args_dtype)
 
@@ -585,7 +585,7 @@ def int_to_float_fn(inputs, out_dtype):
         args_dtype_sz = max(_arg.type.numpy_dtype.itemsize for _arg in inputs)
         args_dtype = np.dtype(f"f{args_dtype_sz}")
 
-        @numba_njit
+        @numba_njit(inline="always")
         def inputs_cast(x):
             return x.astype(args_dtype)
 
@@ -594,36 +594,49 @@ def int_to_float_fn(inputs, out_dtype):
 
 @numba_funcify.register(Dot)
 def numba_funcify_Dot(op, node, **kwargs):
-    # Numba's `np.dot` does not support integer dtypes, so we need to cast to
-    # float.
+    # Numba's `np.dot` does not support integer dtypes, so we need to cast to float.
+    x, y = node.inputs
+    [out] = node.outputs
 
-    out_dtype = node.outputs[0].type.numpy_dtype
-    inputs_cast = int_to_float_fn(node.inputs, out_dtype)
+    x_dtype = x.type.dtype
+    y_dtype = y.type.dtype
+    dot_dtype = f"float{max((32, out.type.numpy_dtype.itemsize * 8))}"
+    out_dtype = out.type.dtype
 
-    @numba_njit
-    def dot(x, y):
-        return np.asarray(np.dot(inputs_cast(x), inputs_cast(y))).astype(out_dtype)
+    if x_dtype == dot_dtype and y_dtype == dot_dtype:
+
+        @numba_njit
+        def dot(x, y):
+            return np.asarray(np.dot(x, y))
 
-    return dot
+    elif x_dtype == dot_dtype and y_dtype != dot_dtype:
 
+        @numba_njit
+        def dot(x, y):
+            return np.asarray(np.dot(x, y.astype(dot_dtype)))
 
-@numba_funcify.register(Softplus)
-def numba_funcify_Softplus(op, node, **kwargs):
-    x_dtype = np.dtype(node.inputs[0].dtype)
+    elif x_dtype != dot_dtype and y_dtype == dot_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)
+        @numba_njit
+        def dot(x, y):
+            return np.asarray(np.dot(x.astype(dot_dtype), y))
 
-    return softplus
+    else:
+
+        @numba_njit()
+        def dot(x, y):
+            return np.asarray(np.dot(x.astype(dot_dtype), y.astype(dot_dtype)))
+
+    if out_dtype == dot_dtype:
+        return dot
+
+    else:
+
+        @numba_njit
+        def dot_with_cast(x, y):
+            return dot(x, y).astype(out_dtype)
+
+    return dot_with_cast
 
 
 @numba_funcify.register(Solve)
@@ -689,11 +702,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.3}/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.3}/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.3}/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.3}/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.3}/pytensor/link/numba/dispatch/tensor_basic.py

@@ -117,17 +117,9 @@ def numba_funcify_ARange(op, **kwargs):
 
 @numba_funcify.register(Join)
 def numba_funcify_Join(op, **kwargs):
-    view = op.view
-
-    if view != -1:
-        # TODO: Where (and why) is this `Join.view` even being used?  From a
-        # quick search, the answer appears to be "nowhere", so we should
-        # probably just remove it.
-        raise NotImplementedError("The `view` parameter to `Join` is not supported")
-
     @numba_basic.numba_njit
     def join(axis, *tensors):
-        return np.concatenate(tensors, numba_basic.to_scalar(axis))
+        return np.concatenate(tensors, axis.item())
 
     return join
 

{pytensor-2.31.1 → pytensor-2.31.3}/pytensor/scan/checkpoints.py

@@ -1,6 +1,5 @@
 import pytensor.tensor.basic as ptb
 from pytensor.scan.basic import scan
-from pytensor.tensor.basic import Join
 from pytensor.tensor.math import ceil, eq, neq
 from pytensor.tensor.subtensor import set_subtensor
 
@@ -127,14 +126,12 @@ def scan_checkpoints(
 
     # Pad the sequences if needed
     if padding:
-        # Since padding could be an empty tensor, Join returns a view of s.
-        join = Join(view=0)
         for i, s in enumerate(sequences):
             overshoots_by = s.shape[0] % save_every_N
             overshoots = neq(overshoots_by, 0)
             n = (save_every_N - overshoots_by) * overshoots
             z = ptb.zeros((n, *s.shape[1:]), dtype=s.dtype)
-            sequences[i] = join(0, s, z)
+            sequences[i] = ptb.join(0, s, z)
 
     # Establish the input variables of the outer scan
     o_sequences = [

{pytensor-2.31.1 → pytensor-2.31.3}/pytensor/scan/op.py

@@ -1647,8 +1647,9 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             p=p, i=node_input_storage, o=node_output_storage, n=node, allow_gc=allow_gc
         ):
             r = p(n, [x[0] for x in i], o)
-            for o in node.outputs:
-                compute_map[o][0] = True
+            if compute_map is not None:
+                for o in node.outputs:
+                    compute_map[o][0] = True
             if allow_gc:
                 self.fn.free()
             return r

{pytensor-2.31.1 → pytensor-2.31.3}/pytensor/scan/rewriting.py

@@ -53,6 +53,7 @@ from pytensor.scan.utils import (
 from pytensor.tensor.basic import (
     Alloc,
     AllocEmpty,
+    atleast_Nd,
     get_scalar_constant_value,
 )
 from pytensor.tensor.elemwise import DimShuffle, Elemwise
@@ -1186,8 +1187,8 @@ def while_scan_merge_subtensor_last_element(fgraph, scan_node):
     return subtensor_merge_replacements
 
 
-def _is_default_scan_buffer(x: TensorVariable) -> bool:
-    node = x.owner
+def _is_default_scan_buffer(final_buffer: TensorVariable, taps: int) -> bool:
+    node = final_buffer.owner
 
     if node is None:
         return False
@@ -1200,8 +1201,10 @@ def _is_default_scan_buffer(x: TensorVariable) -> bool:
     ):
         return False
 
-    x, y, *_ = node.inputs
-    if not (x.owner is not None and isinstance(x.owner.op, AllocEmpty)):
+    init_buffer, init_value, *_ = node.inputs
+    if not (
+        init_buffer.owner is not None and isinstance(init_buffer.owner.op, AllocEmpty)
+    ):
         return False
 
     # The value may have been broadcast to fill in the initial taps.
@@ -1218,10 +1221,16 @@ def _is_default_scan_buffer(x: TensorVariable) -> bool:
     #   1. alloc_empty(2 + nsteps)[:2].broadcastable == x.broadcastable
     # But due to laziness we use the slightly more conservative check:
     #   2. alloc_empty(2 + nsteps).broadcastable == x.broadcastable
-    if broadcasted_by(y, x):
-        return False
-
-    return True
+    if taps > 1:
+        return not broadcasted_by(init_value, init_buffer)
+    else:
+        # In this case we know we have alloc_empty(1 + nsteps, ...)[:1].set(init_value)
+        # The first dimension cannot possibly broadcast in the subtensor assignment,
+        # so we exclude it from `broadcasted_by`. To exclude it we squeeze it out,
+        # after adding any other implicit expand_dims. We select into the first entry of
+        # the buffer, to check for potential broadcasting in other dimensions.
+        init_value_ = atleast_Nd(init_value, n=init_buffer.ndim)
+        return not broadcasted_by(init_value_.squeeze(0), init_buffer[0])
 
 
 def scan_save_mem_rewrite(fgraph, node, backend_supports_output_pre_allocation: bool):
@@ -1574,15 +1583,16 @@ def scan_save_mem_rewrite(fgraph, node, backend_supports_output_pre_allocation:
                 # If the memory for this output has been pre-allocated
                 # before going into the scan op (by an alloc node)
                 if idx < op_info.n_mit_sot + op_info.n_sit_sot:
+                    taps = init_l[i]
                     nw_input = nw_inputs[offset + idx]
 
                     # Recreate default buffers with new size
-                    if _is_default_scan_buffer(nw_input):
-                        extra_size = 1 if required_orphan else val - init_l[i]
+                    if _is_default_scan_buffer(nw_input, taps):
+                        extra_size = 1 if required_orphan else val - taps
                         nw_input = expand_empty(nw_input.owner.inputs[1], extra_size)
                     # Otherwise, just trim with a slice
                     else:
-                        stop = init_l[i] if required_orphan else val
+                        stop = taps if required_orphan else val
                         nw_input = nw_input[:stop]
 
                     nw_inputs[offset + idx] = nw_input
@@ -1626,14 +1636,13 @@ def scan_save_mem_rewrite(fgraph, node, backend_supports_output_pre_allocation:
                     # val == 0 means that we want to keep all intermediate
                     # results for that state, including the initial values.
                     if idx < op_info.n_mit_sot + op_info.n_sit_sot:
+                        taps = init_l[op_info.n_mit_mot + idx]
                         in_idx = offset + idx
                         nw_input = nw_inputs[in_idx]
-                        if _is_default_scan_buffer(nw_input):
+                        if _is_default_scan_buffer(nw_input, taps):
                             nw_input = expand_empty(nw_input.owner.inputs[1], nw_steps)
                         else:
-                            # Number of steps in the initial state
-                            init_l_pt = pt.as_tensor(init_l[op_info.n_mit_mot + idx])
-                            nw_input = nw_input[: (init_l_pt + nw_steps)]
+                            nw_input = nw_input[: (taps + nw_steps)]
                         nw_inputs[in_idx] = nw_input
 
                     elif (

{pytensor-2.31.1 → pytensor-2.31.3}/pytensor/sparse/sandbox/sp2.py

@@ -96,7 +96,7 @@ class Binomial(Op):
 
     def __init__(self, format, dtype):
         self.format = format
-        self.dtype = dtype
+        self.dtype = np.dtype(dtype).name
 
     def make_node(self, n, p, shape):
         n = pt.as_tensor_variable(n)

{pytensor-2.31.1 → pytensor-2.31.3}/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)