Split ScanInfo.tap_array into mit-mot, mit-sot, sit-sot parts

aesara/scan/basic.py

@@ -1123,15 +1123,14 @@ def scan(
     # Step 7. Create the Scan Op
     ##
 
-    tap_array = tuple(tuple(v) for v in mit_sot_tap_array) + tuple(
-        (-1,) for x in range(n_sit_sot)
-    )
     if allow_gc is None:
         allow_gc = config.scan__allow_gc
 
     info = ScanInfo(
-        tap_array=tap_array,
         n_seqs=n_seqs,
+        mit_mot_in_slices=(),
+        mit_sot_in_slices=tuple(tuple(v) for v in mit_sot_tap_array),
+        sit_sot_in_slices=tuple((-1,) for x in range(n_sit_sot)),
         n_mit_mot=n_mit_mot,
         n_mit_mot_outs=n_mit_mot_outs,
         mit_mot_out_slices=tuple(tuple(v) for v in mit_mot_out_slices),

aesara/scan/op.py

@@ -47,7 +47,7 @@ import dataclasses
 import logging
 import time
 from collections import OrderedDict
-from itertools import product
+from itertools import chain, product
 from typing import Callable, List, Optional, Union
 
 import numpy as np
@@ -203,11 +203,9 @@ def copy_var_format(var, as_var):
 
 @dataclasses.dataclass(frozen=True)
 class ScanInfo:
-    tap_array: tuple
-    """
-    This is a tuple containing tuples of inner-output lag/lead values for the
-    mit-mots, mit-sots, and ``[-1]`` for each sit-sot.
-    """
+    mit_mot_in_slices: tuple
+    mit_sot_in_slices: tuple
+    sit_sot_in_slices: tuple
     n_seqs: int
     n_mit_mot: int
     n_mit_mot_outs: int
@@ -219,6 +217,10 @@ class ScanInfo:
     n_non_seqs: int
     as_while: bool
 
+    @property
+    def tap_array(self):
+        return self.mit_mot_in_slices + self.mit_sot_in_slices + self.sit_sot_in_slices
+
 
 TensorConstructorType = Callable[[List[bool], Union[str, np.generic]], TensorType]
 
@@ -235,7 +237,7 @@ class ScanMethodsMixin:
         return list_inputs[1 : 1 + self.info.n_seqs]
 
     def inner_mitmot(self, list_inputs):
-        n_taps = sum(len(x) for x in self.info.tap_array[: self.info.n_mit_mot])
+        n_taps = sum(len(x) for x in self.info.mit_mot_in_slices)
         return list_inputs[self.info.n_seqs : self.info.n_seqs + n_taps]
 
     def outer_mitmot(self, list_inputs):
@@ -251,16 +253,15 @@ class ScanMethodsMixin:
         return list_outputs[: self.info.n_mit_mot]
 
     def mitmot_taps(self):
-        return self.info.tap_array[: self.info.n_mit_mot]
+        return self.info.mit_mot_in_slices
 
     def mitmot_out_taps(self):
         return self.info.mit_mot_out_slices[: self.info.n_mit_mot]
 
     def inner_mitsot(self, list_inputs):
-        n_mitmot_taps = sum(len(x) for x in self.info.tap_array[: self.info.n_mit_mot])
-        ntaps_upto_sit_sot = sum(
-            len(x)
-            for x in self.info.tap_array[: (self.info.n_mit_mot + self.info.n_mit_sot)]
+        n_mitmot_taps = sum(len(x) for x in self.info.mit_mot_in_slices)
+        ntaps_upto_sit_sot = n_mitmot_taps + sum(
+            len(x) for x in self.info.mit_sot_in_slices
         )
         return list_inputs[
             self.info.n_seqs + n_mitmot_taps : self.info.n_seqs + ntaps_upto_sit_sot
@@ -280,14 +281,12 @@ class ScanMethodsMixin:
         ]
 
     def mitsot_taps(self):
-        return self.info.tap_array[
-            self.info.n_mit_mot : self.info.n_mit_mot + self.info.n_mit_sot
-        ]
+        return self.info.mit_sot_in_slices
 
     def inner_sitsot(self, list_inputs):
         n_taps_upto_sit_sot = sum(
             len(x)
-            for x in self.info.tap_array[: (self.info.n_mit_mot + self.info.n_mit_sot)]
+            for x in chain(self.info.mit_mot_in_slices, self.info.mit_sot_in_slices)
         )
         offset = self.info.n_seqs + n_taps_upto_sit_sot
         return list_inputs[offset : offset + self.info.n_sit_sot]
@@ -328,7 +327,7 @@ class ScanMethodsMixin:
     def inner_shared(self, list_inputs):
         n_taps_upto_sit_sot = sum(
             len(x)
-            for x in self.info.tap_array[: (self.info.n_mit_mot + self.info.n_mit_sot)]
+            for x in chain(self.info.mit_mot_in_slices, self.info.mit_sot_in_slices)
         )
         offset = self.info.n_seqs + n_taps_upto_sit_sot + self.info.n_sit_sot
         return list_inputs[offset : offset + self.info.n_shared_outs]
@@ -362,7 +361,7 @@ class ScanMethodsMixin:
     def inner_non_seqs(self, list_inputs):
         n_taps_upto_sit_sot = sum(
             len(x)
-            for x in self.info.tap_array[: (self.info.n_mit_mot + self.info.n_mit_sot)]
+            for x in chain(self.info.mit_mot_in_slices, self.info.mit_sot_in_slices)
         )
         offset = (
             self.info.n_seqs
@@ -427,8 +426,14 @@ class ScanMethodsMixin:
             outer_oidx += 0
 
         # Handle mitmots, mitsots and sitsots variables
-        for i in range(len(self.info.tap_array)):
-            nb_input_taps = len(self.info.tap_array[i])
+        for i, tap in enumerate(
+            chain(
+                self.info.mit_mot_in_slices,
+                self.info.mit_sot_in_slices,
+                self.info.sit_sot_in_slices,
+            )
+        ):
+            nb_input_taps = len(tap)
 
             if i < self.info.n_mit_mot:
                 nb_output_taps = len(self.info.mit_mot_out_slices[i])
@@ -758,7 +763,12 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             self.name = "scan_fn"
 
         # Pre-computing some values to speed up perform
-        self.mintaps = [np.min(x) for x in info.tap_array]
+        self.mintaps = [
+            min(x)
+            for x in chain(
+                info.mit_mot_in_slices, info.mit_sot_in_slices, info.sit_sot_in_slices
+            )
+        ]
         self.mintaps += [0 for x in range(info.n_nit_sot)]
         self.seqs_arg_offset = 1 + info.n_seqs
         self.shared_arg_offset = (
@@ -813,7 +823,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             info = self.info
             input_idx = info.n_seqs
             for mitmot_idx in range(info.n_mit_mot):
-                for inp_tap in info.tap_array[mitmot_idx]:
+                for inp_tap in info.mit_mot_in_slices[mitmot_idx]:
                     if inp_tap in info.mit_mot_out_slices[mitmot_idx]:
                         # Figure out the index of the corresponding output
                         output_idx = sum(
@@ -1292,7 +1302,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
 
             input_idx = self.info.n_seqs
             for mitmot_idx in range(self.info.n_mit_mot):
-                for inp_tap in self.info.tap_array[mitmot_idx]:
+                for inp_tap in self.info.mit_mot_in_slices[mitmot_idx]:
                     if inp_tap in self.info.mit_mot_out_slices[mitmot_idx]:
                         inp = self.inputs[input_idx]
 
@@ -1447,7 +1457,14 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
 
             cython_mintaps = np.asarray(self.mintaps, dtype="int32")
 
-            tap_array_len = tuple(len(x) for x in self.info.tap_array)
+            tap_array_len = tuple(
+                len(x)
+                for x in chain(
+                    self.info.mit_mot_in_slices,
+                    self.info.mit_sot_in_slices,
+                    self.info.sit_sot_in_slices,
+                )
+            )
 
             cython_vector_seqs = np.asarray(self.vector_seqs, dtype="int32")
             cython_vector_outs = np.asarray(self.vector_outs, dtype="int32")
@@ -1506,7 +1523,9 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
                         self.info.n_nit_sot,
                         self.info.as_while,
                         cython_mintaps,
-                        self.info.tap_array,
+                        self.info.mit_mot_in_slices
+                        + self.info.mit_sot_in_slices
+                        + self.info.sit_sot_in_slices,
                         tap_array_len,
                         cython_vector_seqs,
                         cython_vector_outs,
@@ -1679,7 +1698,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
         offset = self.nit_sot_arg_offset + info.n_nit_sot
         other_args = inputs[offset:]
         inner_input_storage = self.fn.input_storage
-        nb_mitmot_in = sum(map(len, info.tap_array[: info.n_mit_mot]))
+        nb_mitmot_in = sum(map(len, info.mit_mot_in_slices))
         old_mitmot_input_storage = [None] * nb_mitmot_in
         old_mitmot_input_data = [None] * nb_mitmot_in
         inner_output_storage = self.fn.output_storage
@@ -1688,7 +1707,14 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
         fn = self.fn.fn
         offset = (
             info.n_seqs
-            + sum(map(len, info.tap_array[: self.n_outs]))
+            + sum(
+                len(x)
+                for x in chain(
+                    info.mit_mot_in_slices,
+                    info.mit_sot_in_slices,
+                    info.sit_sot_in_slices,
+                )
+            )
             + info.n_shared_outs
         )
         for idx in range(len(other_args)):
@@ -1710,17 +1736,23 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
                     inner_input_storage[idx].storage[0] = seqs[idx][i]
 
             offset = info.n_seqs
-            for idx in range(self.n_outs):
+            for idx, taps in enumerate(
+                chain(
+                    info.mit_mot_in_slices,
+                    info.mit_sot_in_slices,
+                    info.sit_sot_in_slices,
+                )
+            ):
                 if self.vector_outs[idx]:
-                    for tap in info.tap_array[idx]:
-                        _idx = (pos[idx] + tap) % store_steps[idx]
+                    for t in taps:
+                        _idx = (pos[idx] + t) % store_steps[idx]
                         inner_input_storage[offset].storage[0] = output_storage[idx][0][
                             _idx : _idx + 1
                         ].reshape(())
                         offset += 1
                 else:
-                    for tap in info.tap_array[idx]:
-                        _idx = (pos[idx] + tap) % store_steps[idx]
+                    for t in taps:
+                        _idx = (pos[idx] + t) % store_steps[idx]
                         inner_input_storage[offset].storage[0] = output_storage[idx][0][
                             _idx
                         ]
@@ -1864,11 +1896,11 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             # 5.3 Copy over the values for mit_mot outputs
             mitmot_inp_offset = 0
             mitmot_out_idx = 0
-            for j in range(info.n_mit_mot):
+            for j, taps in enumerate(info.mit_mot_in_slices):
                 for k in info.mit_mot_out_slices[j]:
                     if self.mitmots_preallocated[mitmot_out_idx]:
                         # This output tap has been preallocated.
-                        inp_idx = mitmot_inp_offset + info.tap_array[j].index(k)
+                        inp_idx = mitmot_inp_offset + taps.index(k)
 
                         # Verify whether the input points to the same data as
                         # it did before the execution of the inner function.
@@ -1899,7 +1931,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
 
                     mitmot_out_idx += 1
 
-                mitmot_inp_offset += len(info.tap_array[j])
+                mitmot_inp_offset += len(taps)
 
             # 5.4 Copy over the values for mit_sot/sit_sot outputs
             begin = info.n_mit_mot
@@ -2141,14 +2173,17 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
 
         n_outs = info.n_mit_mot + info.n_mit_sot + info.n_sit_sot
         outs_shape = []
-        for idx in range(n_outs):
-            abs(min(info.tap_array[idx]))
-            for k in info.tap_array[idx]:
+        for idx, taps in enumerate(
+            chain(
+                info.mit_mot_in_slices, info.mit_sot_in_slices, info.sit_sot_in_slices
+            )
+        ):
+            for k in taps:
                 outs_shape += [input_shapes[idx + info.n_seqs + 1][1:]]
                 # if extra_infer_shape:
+                #     mintap = abs(min(taps))
                 #     corresponding_tap = node.inputs[outer_inp_idx][mintap + k]
                 #     out_equivalent[self.inputs[inner_inp_idx]] = corresponding_tap
-                #     inner_inp_idx += 1
             outer_inp_idx += 1
 
         # shared_outs
@@ -2511,7 +2546,14 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
 
             # Check if the pos-th input is associated with one of the
             # recurrent states
-            x_is_state = pos < sum(len(t) for t in info.tap_array)
+            x_is_state = pos < sum(
+                len(t)
+                for t in chain(
+                    info.mit_mot_in_slices,
+                    info.mit_sot_in_slices,
+                    info.sit_sot_in_slices,
+                )
+            )
 
             if x_is_state and len(idxs) > 0:
                 opos = idxs[0]
@@ -2537,14 +2579,16 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             outer_inp_seqs = [x[n_steps - 1 :: -1] for x in inputs[1 : 1 + info.n_seqs]]
         else:
             outer_inp_seqs = [x[::-1] for x in inputs[1 : 1 + info.n_seqs]]
-        for idx in range(info.n_mit_mot + info.n_mit_sot):
-            mintap = np.min(info.tap_array[idx])
+        for idx, taps in enumerate(
+            chain(info.mit_mot_in_slices, info.mit_sot_in_slices)
+        ):
+            mintap = min(taps)
             if idx < info.n_mit_mot:
                 outmaxtap = np.max(self.mitmot_out_taps()[idx])
             else:
                 outmaxtap = 0
             seq = outs[idx]
-            for k in info.tap_array[idx]:
+            for k in taps:
                 if outmaxtap - k != 0:
                     nw_seq = seq[k - mintap : -(outmaxtap - k)][::-1]
                 else:
@@ -2611,7 +2655,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
         n_mitmot_outs = 0
         n_mitmot_inps = 0
 
-        for idx in range(info.n_mit_mot):
+        for idx, taps in enumerate(info.mit_mot_in_slices):
             if isinstance(dC_douts[idx].type, DisconnectedType):
                 out = outs[idx]
                 outer_inp_mitmot.append(at.zeros_like(out))
@@ -2623,14 +2667,14 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             through_shared = False
             disconnected = True
 
-            for jdx in range(len(info.mit_mot_out_slices[idx])):
+            for mit_mot_out_slice in info.mit_mot_out_slices[idx]:
                 inner_inp_mitmot.append(dC_dXts[out_pos])
-                mitmot_inp_taps[idx].append(-info.mit_mot_out_slices[idx][jdx])
+                mitmot_inp_taps[idx].append(-mit_mot_out_slice)
                 n_mitmot_inps += 1
                 out_pos += 1
 
-            for jdx in range(len(info.tap_array[idx])):
-                tap = -info.tap_array[idx][jdx]
+            for tap in taps:
+                tap = -tap
 
                 # Only create a new inner input if there is not already one
                 # associated with this input tap
@@ -2656,29 +2700,27 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
                             idx
                         ].index(tap)
                         replacement = inner_inp_mitmot[-replacement_idx]
-
-                        info.tap_array[idx]
                         new_inner_out_mitmot = clone_replace(
                             new_inner_out_mitmot, replace=[(to_replace, replacement)]
                         )
 
                     inner_out_mitmot.append(new_inner_out_mitmot)
 
-                if not disconnected_dC_dinps_t[ins_pos]:
-                    disconnected = False
+                disconnected &= disconnected_dC_dinps_t[ins_pos]
 
-                for _sh in self.inner_shared(self_inputs):
-                    if _sh in graph_inputs([dC_dinps_t[ins_pos]]):
-                        through_shared = True
+                through_shared = any(
+                    _sh in graph_inputs([dC_dinps_t[ins_pos]])
+                    for _sh in self.inner_shared(self_inputs)
+                )
 
                 ins_pos += 1
                 n_mitmot_outs += 1
-                mitmot_out_taps[idx].append(-info.tap_array[idx][jdx])
+                mitmot_out_taps[idx].append(tap)
 
                 # Only add the tap as a new input tap if needed
                 if tap not in mitmot_inp_taps[idx]:
                     n_mitmot_inps += 1
-                    mitmot_inp_taps[idx].append(-info.tap_array[idx][jdx])
+                    mitmot_inp_taps[idx].append(tap)
 
             if undefined_msg:
                 type_outs.append(undefined_msg)
@@ -2690,14 +2732,13 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
                 type_outs.append("connected")
 
         offset = info.n_mit_mot
-        for idx in range(info.n_mit_sot):
+        for idx, taps in enumerate(info.mit_sot_in_slices):
             if isinstance(dC_douts[idx + offset].type, DisconnectedType):
                 outer_inp_mitmot.append(outs[idx + offset].zeros_like())
             else:
                 outer_inp_mitmot.append(dC_douts[idx + offset][::-1])
             mitmot_inp_taps.append([])
             mitmot_out_taps.append([])
-            idx_tap = idx + info.n_mit_mot
             inner_inp_mitmot.append(dC_dXts[out_pos])
             out_pos += 1
             n_mitmot_inps += 1
@@ -2705,7 +2746,8 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             through_shared = False
             disconnected = True
             mitmot_inp_taps[idx + offset].append(0)
-            for jdx in range(len(info.tap_array[idx_tap])):
+            for tap in taps:
+                tap = -tap
                 inner_inp_mitmot.append(dC_dXtm1s[ins_pos - info.n_seqs])
 
                 if isinstance(dC_dinps_t[ins_pos].type, NullType):
@@ -2718,13 +2760,15 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
                 else:
                     inner_out_mitmot.append(dC_dinps_t[ins_pos])
 
-                mitmot_inp_taps[idx + offset].append(-info.tap_array[idx_tap][jdx])
-                mitmot_out_taps[idx].append(-info.tap_array[idx_tap][jdx])
-                if not disconnected_dC_dinps_t[ins_pos]:
-                    disconnected = False
-                for _sh in self.inner_shared(self_inputs):
-                    if _sh in graph_inputs([dC_dinps_t[ins_pos]]):
-                        through_shared = True
+                mitmot_inp_taps[idx + offset].append(tap)
+                mitmot_out_taps[idx].append(tap)
+
+                disconnected &= disconnected_dC_dinps_t[ins_pos]
+
+                through_shared = any(
+                    _sh in graph_inputs([dC_dinps_t[ins_pos]])
+                    for _sh in self.inner_shared(self_inputs)
+                )
 
                 n_mitmot_inps += 1
                 ins_pos += 1
@@ -2770,9 +2814,10 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             else:
                 inner_out_mitmot.append(dC_dinps_t[ins_pos])
 
-            for _sh in self.inner_shared(self_inputs):
-                if _sh in graph_inputs([dC_dinps_t[ins_pos]]):
-                    through_shared = True
+            through_shared = any(
+                _sh in graph_inputs([dC_dinps_t[ins_pos]])
+                for _sh in self.inner_shared(self_inputs)
+            )
 
             if isinstance(dC_dinps_t[ins_pos].type, NullType):
                 type_outs.append(dC_dinps_t[ins_pos].type.why_null)
@@ -2860,7 +2905,6 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
                 )
 
         n_sitsot_outs = len(outer_inp_sitsot)
-        new_tap_array = mitmot_inp_taps + [[-1] for k in range(n_sitsot_outs)]
 
         outer_inputs = (
             [grad_steps]
@@ -2884,7 +2928,9 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
         out_info = ScanInfo(
             n_seqs=len(outer_inp_seqs),
             n_mit_sot=0,
-            tap_array=tuple(tuple(v) for v in new_tap_array),
+            mit_mot_in_slices=tuple(tuple(v) for v in mitmot_inp_taps),
+            mit_sot_in_slices=(),
+            sit_sot_in_slices=tuple((-1,) for k in range(n_sitsot_outs)),
             n_mit_mot=len(outer_inp_mitmot),
             n_mit_mot_outs=n_mitmot_outs,
             mit_mot_out_slices=tuple(tuple(v) for v in mitmot_out_taps),
@@ -3065,16 +3111,9 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
         # evan point for the number of nit_sot which I think should just be
         # ignored (?)
 
-        new_tap_array = []
-        b = 0
-        e = info.n_mit_mot
-        new_tap_array += info.tap_array[b:e] * 2
-        b = e
-        e += info.n_mit_sot
-        new_tap_array += info.tap_array[b:e] * 2
-        b = e
-        e += info.n_sit_sot
-        new_tap_array += info.tap_array[b:e] * 2
+        new_mit_mot_in_slices = info.mit_mot_in_slices * 2
+        new_mit_sot_in_slices = info.mit_sot_in_slices * 2
+        new_sit_sot_in_slices = info.sit_sot_in_slices * 2
 
         # Sequences ...
         b = 1
@@ -3095,7 +3134,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
         b = e
         e = e + info.n_mit_mot
         ib = ie
-        ie = ie + int(sum(len(x) for x in info.tap_array[: info.n_mit_mot]))
+        ie = ie + int(sum(len(x) for x in info.mit_mot_in_slices))
         clean_eval_points = []
         for inp, evp in zip(inputs[b:e], eval_points[b:e]):
             if evp is not None:
@@ -3110,14 +3149,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
         b = e
         e = e + info.n_mit_sot
         ib = ie
-        ie = ie + int(
-            sum(
-                len(x)
-                for x in info.tap_array[
-                    info.n_mit_mot : info.n_mit_mot + info.n_mit_sot
-                ]
-            )
-        )
+        ie = ie + int(sum(len(x) for x in info.mit_sot_in_slices))
         clean_eval_points = []
         for inp, evp in zip(inputs[b:e], eval_points[b:e]):
             if evp is not None:
@@ -3217,13 +3249,15 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
 
         out_info = ScanInfo(
             n_seqs=info.n_seqs * 2,
+            mit_mot_in_slices=new_mit_mot_in_slices,
+            mit_sot_in_slices=new_mit_sot_in_slices,
+            sit_sot_in_slices=new_sit_sot_in_slices,
             n_mit_sot=info.n_mit_sot * 2,
             n_sit_sot=info.n_sit_sot * 2,
             n_mit_mot=info.n_mit_mot * 2,
             n_nit_sot=info.n_nit_sot * 2,
             n_shared_outs=info.n_shared_outs,
             n_mit_mot_outs=n_mit_mot_outs * 2,
-            tap_array=tuple(tuple(v) for v in new_tap_array),
             mit_mot_out_slices=tuple(tuple(v) for v in info.mit_mot_out_slices) * 2,
             n_non_seqs=len(inner_other),
             as_while=info.as_while,

aesara/scan/opt.py

@@ -2,6 +2,7 @@
 
 import copy
 import dataclasses
+from itertools import chain
 from sys import maxsize
 from typing import Dict, List, Optional, Tuple
 
@@ -82,7 +83,7 @@ def remove_constants_and_unused_inputs_scan(fgraph, node):
     op = node.op
     # We only need to take care of sequences and other arguments
     st = op.n_seqs
-    st += int(sum(len(x) for x in op.tap_array[: (op.n_mit_mot + op.n_mit_sot)]))
+    st += int(sum(len(x) for x in chain(op.mit_mot_in_slices, op.mit_sot_in_slices)))
     st += op.n_sit_sot
     st += op.n_shared_outs
 
@@ -1147,7 +1148,7 @@ def save_mem_new_scan(fgraph, node):
     c_outs = op.n_mit_mot + op.n_mit_sot + op.n_sit_sot + op.n_nit_sot
 
     init_l = [0 for x in range(op.n_mit_mot)]
-    init_l += [abs(min(v)) for v in op.tap_array[op.n_mit_mot :]]
+    init_l += [abs(min(v)) for v in chain(op.mit_sot_in_slices, op.sit_sot_in_slices)]
     init_l += [0 for x in range(op.n_nit_sot)]
     # 2. Check the clients of each output and see for how many steps
     # does scan need to run
@@ -1678,30 +1679,29 @@ class ScanMerge(GlobalOptimizer):
             inner_ins[idx].append(rename(nd.op.inner_seqs(nd.op.inputs), idx))
             outer_ins += rename(nd.op.outer_seqs(nd.inputs), idx)
 
-        tap_array = ()
         mit_mot_out_slices = ()
 
+        mit_mot_in_slices = ()
         for idx, nd in enumerate(nodes):
-            # MitMot
             inner_ins[idx].append(rename(nd.op.inner_mitmot(nd.op.inputs), idx))
             inner_outs[idx].append(nd.op.inner_mitmot_outs(nd.op.outputs))
-            tap_array += nd.op.mitmot_taps()
+            mit_mot_in_slices += nd.op.mitmot_taps()
             mit_mot_out_slices += nd.op.mitmot_out_taps()
             outer_ins += rename(nd.op.outer_mitmot(nd.inputs), idx)
             outer_outs += nd.op.outer_mitmot_outs(nd.outputs)
 
+        mit_sot_in_slices = ()
         for idx, nd in enumerate(nodes):
-            # MitSot
             inner_ins[idx].append(rename(nd.op.inner_mitsot(nd.op.inputs), idx))
             inner_outs[idx].append(nd.op.inner_mitsot_outs(nd.op.outputs))
-            tap_array += nd.op.mitsot_taps()
+            mit_sot_in_slices += nd.op.mitsot_taps()
             outer_ins += rename(nd.op.outer_mitsot(nd.inputs), idx)
             outer_outs += nd.op.outer_mitsot_outs(nd.outputs)
 
+        sit_sot_in_slices = ()
         for idx, nd in enumerate(nodes):
-            # SitSot
             inner_ins[idx].append(rename(nd.op.inner_sitsot(nd.op.inputs), idx))
-            tap_array += tuple((-1,) for x in range(nd.op.n_sit_sot))
+            sit_sot_in_slices += tuple((-1,) for x in range(nd.op.n_sit_sot))
             inner_outs[idx].append(nd.op.inner_sitsot_outs(nd.op.outputs))
             outer_ins += rename(nd.op.outer_sitsot(nd.inputs), idx)
             outer_outs += nd.op.outer_sitsot_outs(nd.outputs)
@@ -1794,7 +1794,9 @@ class ScanMerge(GlobalOptimizer):
                     new_inner_outs += inner_outs[idx][gr_idx]
 
         info = ScanInfo(
-            tap_array=tap_array,
+            mit_mot_in_slices=mit_mot_in_slices,
+            mit_sot_in_slices=mit_sot_in_slices,
+            sit_sot_in_slices=sit_sot_in_slices,
             n_seqs=sum(nd.op.n_seqs for nd in nodes),
             n_mit_mot=sum(nd.op.n_mit_mot for nd in nodes),
             n_mit_mot_outs=sum(nd.op.n_mit_mot_outs for nd in nodes),
@@ -2212,14 +2214,10 @@ def push_out_dot1_scan(fgraph, node):
                     inner_non_seqs = op.inner_non_seqs(op.inputs)
                     outer_non_seqs = op.outer_non_seqs(node.inputs)
 
-                    st = len(op.mitmot_taps()) + len(op.mitsot_taps())
-
                     new_info = dataclasses.replace(
                         op.info,
-                        tap_array=(
-                            op.info.tap_array[: st + idx]
-                            + op.info.tap_array[st + idx + 1 :]
-                        ),
+                        sit_sot_in_slices=op.info.sit_sot_in_slices[:idx]
+                        + op.info.sit_sot_in_slices[idx + 1 :],
                         n_sit_sot=op.info.n_sit_sot - 1,
                         n_nit_sot=op.info.n_nit_sot + 1,
                     )

aesara/scan/utils.py

@@ -4,6 +4,7 @@ import copy
 import dataclasses
 import logging
 from collections import OrderedDict, namedtuple
+from itertools import chain
 from typing import TYPE_CHECKING, Callable, List, Optional, Sequence, Set, Tuple, Union
 from typing import cast as type_cast
 
@@ -348,11 +349,12 @@ class Validator:
 
 
 def scan_can_remove_outs(op, out_idxs):
-    """
-    Looks at all outputs defined by indices ``out_idxs`` and see whom can be
-    removed from the scan op without affecting the rest. Return two lists,
-    the first one with the indices of outs that can be removed, the second
-    with the outputs that can not be removed.
+    """Look at all outputs defined by indices ``out_idxs`` and determines which can be removed.
+
+    Returns
+    -------
+    two lists, the first one with the indices of outs that can be removed, the
+    second with the outputs that can not be removed.
 
     """
     non_removable = [o for i, o in enumerate(op.outputs) if i not in out_idxs]
@@ -360,9 +362,10 @@ def scan_can_remove_outs(op, out_idxs):
 
     out_ins = []
     offset = op.n_seqs
-    lim = op.n_mit_mot + op.n_mit_sot + op.n_sit_sot
-    for idx in range(lim):
-        n_ins = len(op.info.tap_array[idx])
+    for idx, tap in enumerate(
+        chain(op.mit_mot_in_slices, op.mit_sot_in_slices, op.sit_sot_in_slices)
+    ):
+        n_ins = len(tap)
         out_ins += [op.inputs[offset : offset + n_ins]]
         offset += n_ins
     out_ins += [[] for k in range(op.n_nit_sot)]
@@ -398,7 +401,9 @@ def compress_outs(op, not_required, inputs):
     from aesara.scan.op import ScanInfo
 
     info = ScanInfo(
-        tap_array=(),
+        mit_mot_in_slices=(),
+        mit_sot_in_slices=(),
+        sit_sot_in_slices=(),
         n_seqs=op.info.n_seqs,
         n_mit_mot=0,
         n_mit_mot_outs=0,
@@ -428,23 +433,24 @@ def compress_outs(op, not_required, inputs):
             info = dataclasses.replace(
                 info,
                 n_mit_mot=info.n_mit_mot + 1,
-                tap_array=info.tap_array + (tuple(op.tap_array[offset + idx]),),
+                mit_mot_in_slices=info.mit_mot_in_slices + (op.mit_mot_in_slices[idx],),
                 mit_mot_out_slices=info.mit_mot_out_slices
-                + (tuple(op.mit_mot_out_slices[offset + idx]),),
+                + (op.mit_mot_out_slices[idx],),
             )
             # input taps
-            for jdx in op.tap_array[offset + idx]:
+            for jdx in op.mit_mot_in_slices[idx]:
                 op_inputs += [op.inputs[i_offset]]
                 i_offset += 1
             # output taps
-            for jdx in op.mit_mot_out_slices[offset + idx]:
+            for jdx in op.mit_mot_out_slices[idx]:
                 op_outputs += [op.outputs[o_offset]]
                 o_offset += 1
             # node inputs
             node_inputs += [inputs[ni_offset + idx]]
         else:
-            o_offset += len(op.mit_mot_out_slices[offset + idx])
-            i_offset += len(op.tap_array[offset + idx])
+            o_offset += len(op.mit_mot_out_slices[idx])
+            i_offset += len(op.mit_mot_in_slices[idx])
+
     info = dataclasses.replace(info, n_mit_mot_outs=len(op_outputs))
     offset += op.n_mit_mot
     ni_offset += op.n_mit_mot
@@ -456,10 +462,10 @@ def compress_outs(op, not_required, inputs):
             info = dataclasses.replace(
                 info,
                 n_mit_sot=info.n_mit_sot + 1,
-                tap_array=info.tap_array + (tuple(op.tap_array[offset + idx]),),
+                mit_sot_in_slices=info.mit_sot_in_slices + (op.mit_sot_in_slices[idx],),
             )
             # input taps
-            for jdx in op.tap_array[offset + idx]:
+            for jdx in op.mit_sot_in_slices[idx]:
                 op_inputs += [op.inputs[i_offset]]
                 i_offset += 1
             # output taps
@@ -469,7 +475,7 @@ def compress_outs(op, not_required, inputs):
             node_inputs += [inputs[ni_offset + idx]]
         else:
             o_offset += 1
-            i_offset += len(op.tap_array[offset + idx])
+            i_offset += len(op.mit_sot_in_slices[idx])
 
     offset += op.n_mit_sot
     ni_offset += op.n_mit_sot
@@ -480,7 +486,7 @@ def compress_outs(op, not_required, inputs):
             info = dataclasses.replace(
                 info,
                 n_sit_sot=info.n_sit_sot + 1,
-                tap_array=info.tap_array + (tuple(op.tap_array[offset + idx]),),
+                sit_sot_in_slices=info.sit_sot_in_slices + (op.sit_sot_in_slices[idx],),
             )
             # input taps
             op_inputs += [op.inputs[i_offset]]
@@ -613,8 +619,9 @@ class ScanArgs:
         n_mit_mot = info.n_mit_mot
         n_mit_sot = info.n_mit_sot
 
-        self.mit_mot_in_slices = list(info.tap_array[:n_mit_mot])
-        self.mit_sot_in_slices = list(info.tap_array[n_mit_mot : n_mit_mot + n_mit_sot])
+        self.mit_mot_in_slices = info.mit_mot_in_slices
+        self.mit_sot_in_slices = info.mit_sot_in_slices
+        self.sit_sot_in_slices = info.sit_sot_in_slices
 
         n_mit_mot_ins = sum(len(s) for s in self.mit_mot_in_slices)
         n_mit_sot_ins = sum(len(s) for s in self.mit_sot_in_slices)
@@ -800,14 +807,12 @@ class ScanArgs:
         from aesara.scan.op import ScanInfo
 
         return ScanInfo(
+            mit_mot_in_slices=tuple(tuple(v) for v in self.mit_mot_in_slices),
+            mit_sot_in_slices=tuple(tuple(v) for v in self.mit_sot_in_slices),
+            sit_sot_in_slices=((-1,),) * len(self.inner_in_sit_sot),
             n_seqs=len(self.outer_in_seqs),
             n_mit_mot=len(self.outer_in_mit_mot),
             n_mit_sot=len(self.outer_in_mit_sot),
-            tap_array=(
-                tuple(tuple(v) for v in self.mit_mot_in_slices)
-                + tuple(tuple(v) for v in self.mit_sot_in_slices)
-                + ((-1,),) * len(self.inner_in_sit_sot)
-            ),
             n_sit_sot=len(self.outer_in_sit_sot),
             n_nit_sot=len(self.outer_in_nit_sot),
             n_shared_outs=len(self.outer_in_shared),