Rename n_untraced_sit_sot_outs

pytensor/scan/basic.py

@@ -695,7 +695,7 @@ def scan(
     sit_sot_inner_outputs = []
     sit_sot_rightOrder = []
 
-    n_untraced_sit_sot_outs = 0
+    n_untraced_sit_sot = 0
     untraced_sit_sot_scan_inputs = []
     untraced_sit_sot_inner_inputs = []
     untraced_sit_sot_inner_outputs = []
@@ -763,7 +763,7 @@ def scan(
                     )
                 untraced_sit_sot_scan_inputs.append(actual_arg)
                 untraced_sit_sot_inner_inputs.append(arg)
-                n_untraced_sit_sot_outs += 1
+                n_untraced_sit_sot += 1
                 untraced_sit_sot_rightOrder.append(i)
 
         elif init_out.get("taps", None):
@@ -839,7 +839,7 @@ def scan(
         else:
             _ordered_args[sit_sot_rightOrder[idx]] = [sit_sot_inner_inputs[idx]]
 
-    for idx in range(n_untraced_sit_sot_outs):
+    for idx in range(n_untraced_sit_sot):
         _ordered_args[untraced_sit_sot_rightOrder[idx]] = [
             untraced_sit_sot_inner_inputs[idx]
         ]
@@ -1026,7 +1026,7 @@ def scan(
                 untraced_sit_sot_inner_inputs.append(new_var)
                 untraced_sit_sot_scan_inputs.append(input.variable)
                 untraced_sit_sot_inner_outputs.append(input.update)
-                n_untraced_sit_sot_outs += 1
+                n_untraced_sit_sot += 1
         else:
             no_update_shared_inputs.append(input)
 
@@ -1121,7 +1121,7 @@ def scan(
         mit_mot_out_slices=tuple(tuple(v) for v in mit_mot_out_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_untraced_sit_sot_outs=n_untraced_sit_sot_outs,
+        n_untraced_sit_sot=n_untraced_sit_sot,
         n_nit_sot=n_nit_sot,
         n_non_seqs=len(other_shared_inner_args) + len(other_inner_args),
         as_while=as_while,
@@ -1195,14 +1195,12 @@ def scan(
     offset += n_nit_sot
 
     # Legacy support for explicit untraced sit_sot and those built with update dictionary
-    # Switch to n_untraced_sit_sot_outs after deprecation period
-    n_explicit_untraced_sit_sot_outs = len(untraced_sit_sot_rightOrder)
-    untraced_sit_sot_outs = scan_outs[
-        offset : offset + n_explicit_untraced_sit_sot_outs
-    ]
+    # Switch to n_untraced_sit_sot after deprecation period
+    n_explicit_untraced_sit_sot = len(untraced_sit_sot_rightOrder)
+    untraced_sit_sot_outs = scan_outs[offset : offset + n_explicit_untraced_sit_sot]
 
     # Legacy support: map shared outputs to their updates
-    offset += n_explicit_untraced_sit_sot_outs
+    offset += n_explicit_untraced_sit_sot
     for idx, update_rule in enumerate(scan_outs[offset:]):
         update_map[untraced_sit_sot_scan_inputs[idx]] = update_rule
 

pytensor/scan/op.py

@@ -217,18 +217,18 @@ class ScanInfo:
     mit_sot_in_slices: tuple
     sit_sot_in_slices: tuple
     n_nit_sot: int
-    n_untraced_sit_sot_outs: int
+    n_untraced_sit_sot: int
     n_non_seqs: int
     as_while: bool
 
     @property
     def n_shared_outs(self):
         warnings.warn(
-            "The 'n_shared_outs' property is deprecated. Use 'n_untraced_sit_sot_outs' instead.",
+            "The 'n_shared_outs' property is deprecated. Use 'n_untraced_sit_sot' instead.",
             DeprecationWarning,
             stacklevel=2,
         )
-        return self.n_untraced_sit_sot_outs
+        return self.n_untraced_sit_sot
 
     @property
     def n_mit_mot(self):
@@ -257,7 +257,7 @@ class ScanInfo:
             + sum(len(x) for x in self.mit_mot_in_slices)
             + sum(len(x) for x in self.mit_sot_in_slices)
             + self.n_sit_sot
-            + self.n_untraced_sit_sot_outs
+            + self.n_untraced_sit_sot
             + self.n_non_seqs
         )
 
@@ -268,7 +268,7 @@ class ScanInfo:
             + self.n_mit_sot
             + self.n_sit_sot
             + self.n_nit_sot
-            + self.n_untraced_sit_sot_outs
+            + self.n_untraced_sit_sot
             + int(self.as_while)
         )
 
@@ -281,7 +281,7 @@ class ScanInfo:
             + self.n_mit_sot
             + self.n_sit_sot
             + self.n_nit_sot
-            + self.n_untraced_sit_sot_outs
+            + self.n_untraced_sit_sot
             + self.n_non_seqs
         )
 
@@ -292,7 +292,7 @@ class ScanInfo:
             + self.n_mit_sot
             + self.n_sit_sot
             + self.n_nit_sot
-            + self.n_untraced_sit_sot_outs
+            + self.n_untraced_sit_sot
         )
 
     @property
@@ -419,7 +419,7 @@ class ScanMethodsMixin:
             + self.info.n_mit_mot
             + self.info.n_mit_sot
             + self.info.n_sit_sot
-            + self.info.n_untraced_sit_sot_outs
+            + self.info.n_untraced_sit_sot
         )
         return list_inputs[offset : offset + self.info.n_nit_sot]
 
@@ -438,7 +438,7 @@ class ScanMethodsMixin:
             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_untraced_sit_sot_outs]
+        return list_inputs[offset : offset + self.info.n_untraced_sit_sot]
 
     def inner_shared(self, list_inputs):
         warnings.warn(
@@ -456,7 +456,7 @@ class ScanMethodsMixin:
             + self.info.n_mit_sot
             + self.info.n_sit_sot
         )
-        return list_inputs[offset : offset + self.info.n_untraced_sit_sot_outs]
+        return list_inputs[offset : offset + self.info.n_untraced_sit_sot]
 
     def outer_shared(self, list_inputs):
         warnings.warn(
@@ -471,7 +471,7 @@ class ScanMethodsMixin:
         offset = (
             self.info.n_mit_sot + n_taps + self.info.n_sit_sot + self.info.n_nit_sot
         )
-        return list_outputs[offset : offset + self.info.n_untraced_sit_sot_outs]
+        return list_outputs[offset : offset + self.info.n_untraced_sit_sot]
 
     def inner_shared_outs(self, list_outputs):
         warnings.warn(
@@ -488,7 +488,7 @@ class ScanMethodsMixin:
             + self.info.n_sit_sot
             + self.info.n_nit_sot
         )
-        return list_outputs[offset : offset + self.info.n_untraced_sit_sot_outs]
+        return list_outputs[offset : offset + self.info.n_untraced_sit_sot]
 
     def outer_shared_outs(self, list_outputs):
         warnings.warn(
@@ -507,7 +507,7 @@ class ScanMethodsMixin:
             self.info.n_seqs
             + n_taps_upto_sit_sot
             + self.info.n_sit_sot
-            + self.info.n_untraced_sit_sot_outs
+            + self.info.n_untraced_sit_sot
         )
         return list_inputs[offset:]
 
@@ -519,7 +519,7 @@ class ScanMethodsMixin:
             + self.info.n_mit_sot
             + self.info.n_sit_sot
             + self.info.n_nit_sot
-            + self.info.n_untraced_sit_sot_outs
+            + self.info.n_untraced_sit_sot
         )
         return list_inputs[offset:]
 
@@ -596,7 +596,7 @@ class ScanMethodsMixin:
 
         # This is needed because, for outer inputs (and for outer inputs only)
         # nitsots come *after* untraced_sitsot variables.
-        outer_iidx += self.info.n_untraced_sit_sot_outs
+        outer_iidx += self.info.n_untraced_sit_sot
 
         # Handle nitsots variables
         for i in range(self.info.n_nit_sot):
@@ -612,10 +612,10 @@ class ScanMethodsMixin:
 
         # This is needed because, for outer inputs (and for outer inputs only)
         # nitsots come *after* untraced_sit_sot variables.
-        outer_iidx -= self.info.n_untraced_sit_sot_outs + self.info.n_nit_sot
+        outer_iidx -= self.info.n_untraced_sit_sot + self.info.n_nit_sot
 
         # Handle untraced_sitsot states
-        for i in range(self.info.n_untraced_sit_sot_outs):
+        for i in range(self.info.n_untraced_sit_sot):
             outer_input_indices.append(outer_iidx)
             inner_input_indices.append([inner_iidx])
             inner_output_indices.append([inner_oidx])
@@ -910,7 +910,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             self.seqs_arg_offset + info.n_mit_mot + info.n_mit_sot + info.n_sit_sot
         )
         self.nit_sot_arg_offset = (
-            self.untraced_sit_sot_arg_offset + info.n_untraced_sit_sot_outs
+            self.untraced_sit_sot_arg_offset + info.n_untraced_sit_sot
         )
         # Note: This doesn't include `info.n_nit_sot`s, so it's really a count
         # of the number of outputs generated by taps with inputs
@@ -1647,7 +1647,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
 
                 try:
                     t_fn, n_steps = scan_perform_ext.perform(
-                        self.info.n_untraced_sit_sot_outs,
+                        self.info.n_untraced_sit_sot,
                         self.info.n_mit_mot_outs,
                         self.info.n_seqs,
                         self.info.n_mit_mot,
@@ -1846,7 +1846,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
                     info.sit_sot_in_slices,
                 )
             )
-            + info.n_untraced_sit_sot_outs
+            + info.n_untraced_sit_sot
         )
         for idx in range(len(other_args)):
             inner_input_storage[idx + offset].storage[0] = other_args[idx]
@@ -1892,11 +1892,11 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             a_offset = self.untraced_sit_sot_arg_offset
             o_offset = self.n_outs + info.n_nit_sot
             if i == 0:
-                for j in range(info.n_untraced_sit_sot_outs):
+                for j in range(info.n_untraced_sit_sot):
                     inner_input_storage[offset].storage[0] = inputs[a_offset + j]
                     offset += 1
             else:
-                for j in range(info.n_untraced_sit_sot_outs):
+                for j in range(info.n_untraced_sit_sot):
                     inner_input_storage[offset].storage[0] = output_storage[
                         o_offset + j
                     ][0]
@@ -1930,12 +1930,12 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
 
             # 4.3. Collect slices for untraced sitsot outputs
             offset += self.n_outs + info.n_nit_sot - info.n_mit_mot
-            for idx in range(info.n_untraced_sit_sot_outs):
+            for idx in range(info.n_untraced_sit_sot):
                 inner_output_storage[idx + offset].storage[0] = None
 
             # 4.4. If there is a condition add it to the mix
             if info.as_while:
-                pdx = offset + info.n_untraced_sit_sot_outs
+                pdx = offset + info.n_untraced_sit_sot
                 inner_output_storage[pdx].storage[0] = None
 
             # 4.5. Keep a reference to the variables (ndarrays,
@@ -2004,7 +2004,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
 
             dt_fn = time.perf_counter() - t0_fn
             if info.as_while:
-                pdx = offset + info.n_untraced_sit_sot_outs
+                pdx = offset + info.n_untraced_sit_sot
                 cond = inner_output_storage[pdx].storage[0] == 0
 
             t_fn += dt_fn
@@ -2151,7 +2151,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             # 5.6 Copy over the values for outputs corresponding to untraced sitsot
             # variables
             begin = end
-            end += info.n_untraced_sit_sot_outs
+            end += info.n_untraced_sit_sot
             for j in range(begin, end):
                 jout = j + offset_out
                 output_storage[j][0] = inner_output_storage[jout].storage[0]
@@ -2301,11 +2301,11 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
 
         # untraced sit_sot outputs
         offset = 1 + info.n_seqs + n_outs
-        for idx in range(info.n_untraced_sit_sot_outs):
+        for idx in range(info.n_untraced_sit_sot):
             outs_shape += [input_shapes[idx + offset]]
 
         # non_sequences
-        offset += info.n_nit_sot + info.n_untraced_sit_sot_outs
+        offset += info.n_nit_sot + info.n_untraced_sit_sot
         inner_ins_shapes = seqs_shape + outs_shape + input_shapes[offset:]
         assert len(inner_ins_shapes) == len(self.inner_inputs)
 
@@ -2347,7 +2347,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
                 # in the inner function.
                 r = node.outputs[n_outs + x]
                 assert r.ndim == 1 + len(out_shape_x)
-                shp = [node.inputs[offset + info.n_untraced_sit_sot_outs + x]]
+                shp = [node.inputs[offset + info.n_untraced_sit_sot + x]]
                 for i, shp_i in zip(range(1, r.ndim), out_shape_x, strict=True):
                     # Validate shp_i. v_shape_i is either None (if invalid),
                     # or a (variable, Boolean) tuple. The Boolean indicates
@@ -2364,7 +2364,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
                         shp.append(v_shp_i[0])
                 scan_outs.append(tuple(shp))
 
-        scan_outs += list(input_shapes[offset : offset + info.n_untraced_sit_sot_outs])
+        scan_outs += list(input_shapes[offset : offset + info.n_untraced_sit_sot])
         # if we are dealing with a repeat-until, then we do not know the
         # leading dimension so we replace it for every entry with Shape_i
         if info.as_while:
@@ -2437,8 +2437,6 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
         return connection_pattern
 
     def L_op(self, inputs, outs, dC_douts):
-        if not isinstance(outs, list | tuple):
-            outs = [outs]
         # `grad_step` equals the number of steps the original scan node has
         # done (if the original scan is a while loop than this number is the
         # length of the output sequence)
@@ -2695,7 +2693,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
                 pass
             elif isinstance(dC_dXtm1.type, NullType):
                 # The new gradient is undefined, this makes the accumulated
-                # gradient undefined as weell
+                # gradient undefined as well
                 dC_dinps_t[dx + info.n_seqs] = dC_dXtm1
             else:
                 dC_dinps_t[dx + info.n_seqs] += dC_dXtm1
@@ -3062,7 +3060,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             mit_sot_in_slices=(),
             sit_sot_in_slices=tuple((-1,) for k in range(n_sitsot_outs)),
             n_nit_sot=n_nit_sot,
-            n_untraced_sit_sot_outs=0,
+            n_untraced_sit_sot=0,
             n_non_seqs=len(self.outer_untraced_sit_sot(inputs))
             + len(self.outer_non_seqs(inputs)),
             as_while=False,
@@ -3149,7 +3147,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
 
         start = len(gradients)
         node = outs[0].owner
-        for idx in range(info.n_untraced_sit_sot_outs):
+        for idx in range(info.n_untraced_sit_sot):
             disconnected = True
             connected_flags = self.connection_pattern(node)[idx + start]
             for dC_dout, connected in zip(dC_douts, connected_flags, strict=True):
@@ -3211,7 +3209,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
         self_inputs = self.inner_inputs
         rop_of_inputs = (
             self_inputs[: info.n_seqs + self.n_outs]
-            + self_inputs[info.n_seqs + self.n_outs + info.n_untraced_sit_sot_outs :]
+            + self_inputs[info.n_seqs + self.n_outs + info.n_untraced_sit_sot :]
         )
         self_outputs = self.inner_outputs
 
@@ -3221,8 +3219,8 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             rop_self_outputs = self_outputs[:-1]
         else:
             rop_self_outputs = self_outputs
-        if info.n_untraced_sit_sot_outs > 0:
-            rop_self_outputs = rop_self_outputs[: -info.n_untraced_sit_sot_outs]
+        if info.n_untraced_sit_sot > 0:
+            rop_self_outputs = rop_self_outputs[: -info.n_untraced_sit_sot]
         rop_outs = Rop(
             rop_self_outputs,
             rop_of_inputs,
@@ -3308,9 +3306,9 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
 
         # Untraced outs ...
         b = e
-        e = e + info.n_untraced_sit_sot_outs
+        e = e + info.n_untraced_sit_sot
         ib = ie
-        ie = ie + info.n_untraced_sit_sot_outs
+        ie = ie + info.n_untraced_sit_sot
         scan_untraced = inputs[b:e]
         inner_untraced = self_inputs[ib:ie]
 
@@ -3346,7 +3344,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
         e = e + info.n_nit_sot
         inner_out_nit_sot = self_outputs[b:e] + rop_outs[b:e]
         b = e
-        e = e + info.n_untraced_sit_sot_outs
+        e = e + info.n_untraced_sit_sot
         inner_out_untraced = self_outputs[b:e]
 
         inner_ins = (
@@ -3385,7 +3383,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
             mit_sot_in_slices=new_mit_sot_in_slices,
             sit_sot_in_slices=new_sit_sot_in_slices,
             n_nit_sot=info.n_nit_sot * 2,
-            n_untraced_sit_sot_outs=info.n_untraced_sit_sot_outs,
+            n_untraced_sit_sot=info.n_untraced_sit_sot,
             n_non_seqs=len(inner_other),
             as_while=info.as_while,
         )
@@ -3417,7 +3415,7 @@ class Scan(Op, ScanMethodsMixin, HasInnerGraph):
         b = e + info.n_nit_sot
         e = e + info.n_nit_sot * 2
         final_outs += outputs[b:e]
-        final_outs += [None] * info.n_untraced_sit_sot_outs
+        final_outs += [None] * info.n_untraced_sit_sot
 
         return final_outs
 

pytensor/scan/rewriting.py

@@ -110,7 +110,7 @@ def remove_constants_and_unused_inputs_scan(fgraph, node):
         sum(len(x) for x in chain(op_info.mit_mot_in_slices, op_info.mit_sot_in_slices))
     )
     st += op_info.n_sit_sot
-    st += op_info.n_untraced_sit_sot_outs
+    st += op_info.n_untraced_sit_sot
 
     op_ins = op.inner_inputs
     op_outs = op.inner_outputs
@@ -126,7 +126,7 @@ def remove_constants_and_unused_inputs_scan(fgraph, node):
         + op_info.n_mit_sot
         + op_info.n_sit_sot
         + op_info.n_nit_sot
-        + op_info.n_untraced_sit_sot_outs
+        + op_info.n_untraced_sit_sot
         + 1
     )
     outer_non_seqs = node.inputs[st:]
@@ -1628,7 +1628,7 @@ def scan_save_mem_rewrite(fgraph, node, backend_supports_output_pre_allocation:
                         + idx
                         + op_info.n_seqs
                         + 1
-                        + op_info.n_untraced_sit_sot_outs
+                        + op_info.n_untraced_sit_sot
                     )
                     if nw_inputs[pos] == node.inputs[0]:
                         nw_inputs[pos] = 1 if required_orphan else val
@@ -1662,7 +1662,7 @@ def scan_save_mem_rewrite(fgraph, node, backend_supports_output_pre_allocation:
                     elif (
                         idx < op_info.n_mit_sot + op_info.n_sit_sot + op_info.n_nit_sot
                     ):
-                        in_idx = offset + idx + op_info.n_untraced_sit_sot_outs
+                        in_idx = offset + idx + op_info.n_untraced_sit_sot
                         if nw_inputs[in_idx] == node.inputs[0]:
                             nw_inputs[in_idx] = nw_steps
 
@@ -1980,9 +1980,7 @@ class ScanMerge(GraphRewriter):
             mit_sot_in_slices=mit_sot_in_slices,
             sit_sot_in_slices=sit_sot_in_slices,
             n_nit_sot=sum(nd.op.info.n_nit_sot for nd in nodes),
-            n_untraced_sit_sot_outs=sum(
-                nd.op.info.n_untraced_sit_sot_outs for nd in nodes
-            ),
+            n_untraced_sit_sot=sum(nd.op.info.n_untraced_sit_sot for nd in nodes),
             n_non_seqs=n_non_seqs,
             as_while=as_while,
         )

pytensor/scan/utils.py

@@ -371,7 +371,7 @@ def scan_can_remove_outs(op, out_idxs):
         offset += n_ins
     out_ins += [[] for k in range(op.info.n_nit_sot)]
     out_ins += [
-        [op.inner_inputs[offset + k]] for k in range(op.info.n_untraced_sit_sot_outs)
+        [op.inner_inputs[offset + k]] for k in range(op.info.n_untraced_sit_sot)
     ]
 
     added = True
@@ -411,7 +411,7 @@ def compress_outs(op, not_required, inputs):
         mit_sot_in_slices=(),
         sit_sot_in_slices=(),
         n_nit_sot=0,
-        n_untraced_sit_sot_outs=0,
+        n_untraced_sit_sot=0,
         n_non_seqs=0,
         as_while=op_info.as_while,
     )
@@ -517,18 +517,18 @@ def compress_outs(op, not_required, inputs):
             info = dataclasses.replace(info, n_nit_sot=info.n_nit_sot + 1)
             op_outputs += [op.inner_outputs[o_offset]]
             o_offset += 1
-            nit_sot_ins += [inputs[ni_offset + idx + op_info.n_untraced_sit_sot_outs]]
+            nit_sot_ins += [inputs[ni_offset + idx + op_info.n_untraced_sit_sot]]
         else:
             o_offset += 1
 
     offset += op_info.n_nit_sot
     shared_ins = []
-    for idx in range(op_info.n_untraced_sit_sot_outs):
+    for idx in range(op_info.n_untraced_sit_sot):
         if offset + idx not in not_required:
             map_old_new[offset + idx] = curr_pos
             curr_pos += 1
             info = dataclasses.replace(
-                info, n_untraced_sit_sot_outs=info.n_untraced_sit_sot_outs + 1
+                info, n_untraced_sit_sot=info.n_untraced_sit_sot + 1
             )
             op_outputs += [op.inner_outputs[o_offset]]
             o_offset += 1
@@ -543,9 +543,7 @@ def compress_outs(op, not_required, inputs):
     # other stuff
     op_inputs += op.inner_inputs[i_offset:]
     info = dataclasses.replace(info, n_non_seqs=len(op.inner_inputs[i_offset:]))
-    node_inputs += inputs[
-        ni_offset + op_info.n_untraced_sit_sot_outs + op_info.n_nit_sot :
-    ]
+    node_inputs += inputs[ni_offset + op_info.n_untraced_sit_sot + op_info.n_nit_sot :]
     if op_info.as_while:
         op_outputs += [op.inner_outputs[o_offset]]
         map_old_new[o_offset] = len(op_outputs) - 1
@@ -664,11 +662,11 @@ class ScanArgs:
         p += n_sit_sot
         q += n_sit_sot
 
-        n_untraced_sit_sot_outs = info.n_untraced_sit_sot_outs
-        self.outer_in_shared = list(outer_inputs[p : p + n_untraced_sit_sot_outs])
-        self.inner_in_shared = list(inner_inputs[q : q + n_untraced_sit_sot_outs])
-        p += n_untraced_sit_sot_outs
-        q += n_untraced_sit_sot_outs
+        n_untraced_sit_sot = info.n_untraced_sit_sot
+        self.outer_in_shared = list(outer_inputs[p : p + n_untraced_sit_sot])
+        self.inner_in_shared = list(inner_inputs[q : q + n_untraced_sit_sot])
+        p += n_untraced_sit_sot
+        q += n_untraced_sit_sot
 
         n_nit_sot = info.n_nit_sot
         self.outer_in_nit_sot = list(outer_inputs[p : p + n_nit_sot])
@@ -708,10 +706,10 @@ class ScanArgs:
         p += n_nit_sot
         q += n_nit_sot
 
-        self.outer_out_shared = list(outer_outputs[p : p + n_untraced_sit_sot_outs])
-        self.inner_out_shared = list(inner_outputs[q : q + n_untraced_sit_sot_outs])
-        p += n_untraced_sit_sot_outs
-        q += n_untraced_sit_sot_outs
+        self.outer_out_shared = list(outer_outputs[p : p + n_untraced_sit_sot])
+        self.inner_out_shared = list(inner_outputs[q : q + n_untraced_sit_sot])
+        p += n_untraced_sit_sot
+        q += n_untraced_sit_sot
 
         assert p == len(outer_outputs)
         assert q == len(inner_outputs)
@@ -822,7 +820,7 @@ class ScanArgs:
             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_nit_sot=len(self.outer_in_nit_sot),
-            n_untraced_sit_sot_outs=len(self.outer_in_shared),
+            n_untraced_sit_sot=len(self.outer_in_shared),
             n_non_seqs=len(self.inner_in_non_seqs),
             as_while=self.as_while,
         )

tests/link/jax/test_scan.py

@@ -651,7 +651,7 @@ def test_trace_truncation_regression_bug():
             mit_sot_in_slices=(),
             sit_sot_in_slices=((-1,),),
             n_nit_sot=0,
-            n_untraced_sit_sot_outs=0,
+            n_untraced_sit_sot=0,
             n_non_seqs=0,
             as_while=False,
         ),

tests/link/numba/test_scan.py

@@ -86,7 +86,7 @@ from tests.scan.test_basic import ScanCompatibilityTests
             3,
             [],
             [np.array([0.50100236, 2.16822932, 1.36326596])],
-            lambda op: op.info.n_untraced_sit_sot_outs > 0,
+            lambda op: op.info.n_untraced_sit_sot > 0,
         ),
         # mit-sot (that's also a type of sit-sot)
         (

tests/scan/test_basic.py

@@ -4095,7 +4095,7 @@ class TestExamples:
             [{}],
             [],
             3,
-            lambda op: op.info.n_untraced_sit_sot_outs > 0,
+            lambda op: op.info.n_untraced_sit_sot > 0,
         ),
         # mit-sot (that's also a type of sit-sot)
         (
@@ -4292,7 +4292,7 @@ def test_scan_mode_compatibility(scan_mode):
         mit_sot_in_slices=(),
         sit_sot_in_slices=(),
         n_nit_sot=0,
-        n_untraced_sit_sot_outs=0,
+        n_untraced_sit_sot=0,
         n_non_seqs=0,
         as_while=False,
     )