Allow non-TensorVariable types to be traced in new Scan Op

Author: ricardoV94

pytensor/loop/op.py

@@ -7,18 +7,12 @@
 from pytensor.graph import Apply, FunctionGraph, Op, Type, node_rewriter
 from pytensor.graph.rewriting.basic import in2out
 from pytensor.scalar import constant
-from pytensor.tensor import (
-    NoneConst,
-    add,
-    and_,
-    empty,
-    get_scalar_constant_value,
-    set_subtensor,
-)
+from pytensor.tensor import add, and_, empty, get_scalar_constant_value, set_subtensor
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.shape import Shape_i
 from pytensor.tensor.type import DenseTensorType, TensorType
 from pytensor.tensor.type_other import NoneTypeT
+from pytensor.typed_list import TypedListType, append, make_empty_list
 
 
 def validate_loop_update_types(update):
@@ -176,8 +170,7 @@ def __init__(
                     )
                 )
             else:
-                # We can't concatenate all types of states, such as RandomTypes
-                self.trace_types.append(NoneConst.type)
+                self.trace_types.append(TypedListType(state_type))
 
         self.constant_types = [inp.type for inp in update_fg.inputs[self.n_states :]]
         self.n_constants = len(self.constant_types)
@@ -312,10 +305,6 @@ def scan(fn, idx, initial_states, constants, max_iters):
         if fgraph.clients[trace]
     ]
 
-    # Check that outputs that cannot be converted into sequences (such as RandomTypes) are not being referenced
-    for trace_idx in used_traces_idxs:
-        assert not isinstance(old_states[trace_idx].type, NoneTypeT)
-
     # Inputs to the new Loop
     max_iters = node.inputs[0]
     init_states = node.inputs[1 : 1 + op.n_states]
@@ -324,6 +313,8 @@ def scan(fn, idx, initial_states, constants, max_iters):
             (max_iters, *tuple(init_states[trace_idx].shape)),
             dtype=init_states[trace_idx].dtype,
         )
+        if isinstance(init_states[trace_idx].type, DenseTensorType)
+        else make_empty_list(init_states[trace_idx].type)
         for trace_idx in used_traces_idxs
     ]
     constants = node.inputs[1 + op.n_states :]
@@ -376,6 +367,7 @@ def scan(fn, idx, initial_states, constants, max_iters):
     # Inner traces
     inner_states = update_fg.inputs[: op.n_states]
     inner_traces = [init_trace.type() for init_trace in init_traces]
+
     for s, t in zip(inner_states, inner_traces):
         t.name = "trace"
         if s.name:
@@ -387,6 +379,8 @@ def scan(fn, idx, initial_states, constants, max_iters):
     inner_while_cond, *inner_next_states = update_fg.outputs
     inner_next_traces = [
         set_subtensor(prev_trace[inner_idx], inner_next_states[trace_idx])
+        if isinstance(prev_trace.type, DenseTensorType)
+        else append(prev_trace, inner_next_states[trace_idx])
         for trace_idx, prev_trace in zip(used_traces_idxs, inner_traces)
     ]
     for t in inner_next_traces:
@@ -429,7 +423,7 @@ def scan(fn, idx, initial_states, constants, max_iters):
     replacements = dict(zip(old_states, new_states))
     for trace_idx, new_trace in zip(used_traces_idxs, new_traces):
         # If there is no while condition, the whole trace will be used
-        if op.has_while_condition:
+        if op.has_while_condition and isinstance(new_trace.type, DenseTensorType):
             new_trace = new_trace[:final_idx]
         replacements[old_traces[trace_idx]] = new_trace
     return replacements

tests/loop/test_op.py

@@ -7,8 +7,9 @@
 from pytensor.loop.op import Loop, Scan
 from pytensor.tensor import constant, empty, lscalar, scalar, vector
 from pytensor.tensor.random import normal
+from pytensor.tensor.random.type import RandomGeneratorType
 from pytensor.tensor.subtensor import Subtensor
-from pytensor.tensor.type_other import NoneTypeT
+from pytensor.typed_list import TypedListType
 
 
 def test_loop_basic():
@@ -152,10 +153,16 @@ def test_fori_random_scan():
         [constant(np.array(True)), *normal(rng=rng).owner.outputs[::-1]],
     )
 
-    _, new_rng, ys, rngs = Scan(update_fg=update_fg)(n_iters, dummy_init, rng_shared)
-    assert isinstance(rngs.type, NoneTypeT)
-
-    fn = function([], ys, updates={rng_shared: new_rng})
-
-    np.testing.assert_array_equal(fn(), rng_test.normal(size=5))
-    np.testing.assert_array_equal(fn(), rng_test.normal(size=5))
+    last_y, last_rng, ys, rngs = Scan(update_fg=update_fg)(
+        n_iters, dummy_init, rng_shared
+    )
+    assert isinstance(last_rng.type, RandomGeneratorType)
+    assert isinstance(rngs.type, TypedListType)
+    assert isinstance(rngs.type.ttype, RandomGeneratorType)
+
+    fn = function([], [ys, rngs], updates={rng_shared: last_rng})
+    for i in range(2):
+        ys_res, rngs_res = fn()
+        for y_res, rng_res in zip(ys_res, rngs_res):
+            np.testing.assert_almost_equal(y_res, rng_test.normal())
+            assert rng_res.__getstate__() == rng_test.__getstate__()