Adding workaround so more graphs with Blockwise(Scans) can be

vectorized

Author: zaxtax

pytensor/scan/op.py

@@ -1500,6 +1500,10 @@ def make_thunk(self, node, storage_map, compute_map, no_recycling, impl=None):
         node_input_storage = [storage_map[r] for r in node.inputs]
         node_output_storage = [storage_map[r] for r in node.outputs]
 
+        # HACK: Here to handle Blockwise Scans
+        if compute_map is None:
+            compute_map = {out: [False] for out in node.outputs}
+
         # Analyse the compile inner function to determine which inputs and
         # outputs are on the gpu and speed up some checks during the execution
         outs_is_tensor = [

tests/scan/test_basic.py

@@ -27,6 +27,7 @@
 from pytensor.compile.sharedvalue import shared
 from pytensor.configdefaults import config
 from pytensor.gradient import NullTypeGradError, Rop, disconnected_grad, grad, hessian
+from pytensor.graph import vectorize_graph
 from pytensor.graph.basic import Apply, ancestors, equal_computations
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import Op
@@ -1178,6 +1179,17 @@ def get_sum_of_grad(input0, input1):
 
         utt.verify_grad(get_sum_of_grad, inputs_test_values, rng=rng)
 
+    def test_blockwise_scan(self):
+        x = pt.tensor("x", shape=())
+        out, _ = scan(lambda x: x + 1, outputs_info=[x], n_steps=10)
+        x_vec = pt.tensor("x_vec", shape=(None,))
+        out_vec = vectorize_graph(out, {x: x_vec})
+
+        fn = function([x_vec], out_vec)
+        o1 = fn([1, 2, 3])
+        o2 = np.arange(2, 12) + np.arange(3).reshape(-1, 1)
+        assert np.allclose(o1, o2)
+
     def test_connection_pattern(self):
         """Test `Scan.connection_pattern` in the presence of recurrent outputs with multiple taps."""