pytensor 2.31.2__tar.gz → 2.31.3__tar.gz

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

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

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-05-24T11:34:23+0800",
+ "date": "2025-06-02T15:09:01+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "2d414d4129882861b867aaf0d9abcf3135e3d5eb",
- "version": "2.31.2"
+ "full-revisionid": "ff98ab8f8d2e6b78539a9225bde5e18a49ae6b81",
+ "version": "2.31.3"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.31.2 → 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.2 → 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.2 → pytensor-2.31.3}/pytensor/link/numba/dispatch/basic.py

@@ -565,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)
 
@@ -584,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)
 
@@ -593,17 +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))
+
+    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)))
+
+    elif x_dtype != dot_dtype and y_dtype == dot_dtype:
+
+        @numba_njit
+        def dot(x, y):
+            return np.asarray(np.dot(x.astype(dot_dtype), y))
+
+    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
+    return dot_with_cast
 
 
 @numba_funcify.register(Solve)

{pytensor-2.31.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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)