Start deprecating shared updates API in Scan

Using DeprecationWarning to keep it visible only for devs for now

Author: ricardoV94

pytensor/gradient.py

@@ -2188,7 +2188,7 @@ def hessian(cost, wrt, consider_constant=None, disconnected_inputs="raise"):
         # It is possible that the inputs are disconnected from expr,
         # even if they are connected to cost.
         # This should not be an error.
-        hess, updates = pytensor.scan(
+        hess = pytensor.scan(
             lambda i, y, x: grad(
                 y[i],
                 x,
@@ -2197,9 +2197,7 @@ def hessian(cost, wrt, consider_constant=None, disconnected_inputs="raise"):
             ),
             sequences=pytensor.tensor.arange(expr.shape[0]),
             non_sequences=[expr, input],
-        )
-        assert not updates, (
-            "Scan has returned a list of updates; this should not happen."
+            return_updates=False,
         )
         hessians.append(hess)
     return as_list_or_tuple(using_list, using_tuple, hessians)

pytensor/scan/basic.py

@@ -168,6 +168,26 @@ def isNaN_or_Inf_or_None(x):
     return isNone or isNaN or isInf or isStr
 
 
+def _manage_output_api_change(outputs, updates, return_updates):
+    if return_updates:
+        warnings.warn(
+            "Scan return signature well change. Updates dict will not be returned, only the first argument. "
+            "Pass `return_updates=False` to conform to the new API and avoid this warning",
+            DeprecationWarning,
+            # Only meant for developers for now. Switch to FutureWarning to warn users, before removing.
+            stacklevel=2,
+        )
+    else:
+        if updates:
+            raise ValueError(
+                f"return_updates=False but Scan produced updates {updates}."
+                "Make sure to use outputs_info to handle all recurrent states, and not rely on shared variable updates."
+            )
+        return outputs
+
+    return outputs, updates
+
+
 def scan(
     fn,
     sequences=None,
@@ -182,6 +202,7 @@ def scan(
     allow_gc=None,
     strict=False,
     return_list=False,
+    return_updates: bool = True,
 ):
     r"""This function constructs and applies a `Scan` `Op` to the provided arguments.
 
@@ -878,6 +899,11 @@ def wrap_into_list(x):
         raw_inner_outputs = fn(*args)
 
     condition, outputs, updates = get_updates_and_outputs(raw_inner_outputs)
+    if updates:
+        warnings.warn(
+            "Updates functionality in Scan are deprecated. Use explicit outputs_info and build shared update expressions manually, even for RNGs.",
+            DeprecationWarning,  # Only meant for developers for now. Switch to FutureWarning to warn users, before removing.
+        )
     if condition is not None:
         as_while = True
     else:
@@ -900,7 +926,7 @@ def wrap_into_list(x):
         if not return_list and len(outputs) == 1:
             outputs = outputs[0]
 
-        return (outputs, updates)
+        return _manage_output_api_change(outputs, updates, return_updates)
 
     ##
     # Step 4. Compile the dummy function
@@ -919,6 +945,8 @@ def wrap_into_list(x):
     fake_outputs = clone_replace(
         outputs, replace=dict(zip(non_seqs, fake_nonseqs, strict=True))
     )
+    # TODO: Once we don't treat shared variables specially we should use `truncated_graph_inputs`
+    #  to find implicit inputs in a way that reduces the size of the inner function
     known_inputs = [*args, *fake_nonseqs]
     extra_inputs = [
         x for x in explicit_graph_inputs(fake_outputs) if x not in known_inputs
@@ -1074,7 +1102,7 @@ def wrap_into_list(x):
         if not isinstance(arg, SharedVariable | Constant)
     ]
 
-    inner_replacements.update(dict(zip(other_scan_args, other_inner_args, strict=True)))
+    inner_replacements.update(dict(zip(other_scan_args, other_inner_args, strict=True)))  # type: ignore[arg-type]
 
     if strict:
         non_seqs_set = set(non_sequences if non_sequences is not None else [])
@@ -1123,7 +1151,7 @@ def wrap_into_list(x):
     if condition is not None:
         inner_outs.append(condition)
 
-    new_outs = clone_replace(inner_outs, replace=inner_replacements)
+    new_outs = clone_replace(inner_outs, replace=inner_replacements)  # type: ignore[arg-type]
 
     ##
     # Step 7. Create the Scan Op
@@ -1211,12 +1239,14 @@ def remove_dimensions(outs, offsets=None):
 
     offset += n_nit_sot
 
-    # Support for explicit untraced sit_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
     ]
 
+    # Legacy support: map shared outputs to their updates
     offset += n_explicit_untraced_sit_sot_outs
     for idx, update_rule in enumerate(scan_outs[offset:]):
         update_map[untraced_sit_sot_scan_inputs[idx]] = update_rule
@@ -1245,8 +1275,8 @@ def remove_dimensions(outs, offsets=None):
             update_map[sit_sot_shared[abs(pos) - 1]] = _scan_out_list[idx][-1]
     scan_out_list = [x for x in scan_out_list if x is not None]
     if not return_list and len(scan_out_list) == 1:
-        scan_out_list = scan_out_list[0]
+        scan_out_list = scan_out_list[0]  # type: ignore[assignment]
     elif len(scan_out_list) == 0:
-        scan_out_list = None
+        scan_out_list = None  # type: ignore[assignment]
 
-    return scan_out_list, update_map
+    return _manage_output_api_change(scan_out_list, update_map, return_updates)

pytensor/scan/checkpoints.py

@@ -13,6 +13,7 @@ def scan_checkpoints(
     n_steps=None,
     save_every_N=10,
     padding=True,
+    return_updates=True,
 ):
     """Scan function that uses less memory, but is more restrictive.
 
@@ -157,31 +158,34 @@ def outer_step(*args):
         ] * len(new_nitsots)
 
         # Call the user-provided function with the proper arguments
-        results, updates = scan(
+        results_and_updates = scan(
             fn=fn,
             sequences=i_sequences[:-1],
             outputs_info=i_outputs_infos,
             non_sequences=i_non_sequences,
             name=name + "_inner",
             n_steps=i_sequences[-1],
+            return_updates=return_updates,
         )
+        if return_updates:
+            results, updates = results_and_updates
+        else:
+            results = results_and_updates
+            updates = {}
+
         if not isinstance(results, list):
             results = [results]
 
         # Keep only the last timestep of every output but keep all the updates
-        if not isinstance(results, list):
-            return results[-1], updates
-        else:
-            return [r[-1] for r in results], updates
+        return [r[-1] for r in results], updates
 
-    results, updates = scan(
+    return scan(
         fn=outer_step,
         sequences=o_sequences,
         outputs_info=outputs_info,
         non_sequences=o_nonsequences,
         name=name + "_outer",
         n_steps=o_n_steps,
         allow_gc=True,
+        return_updates=return_updates,
     )
-
-    return results, updates

pytensor/scan/views.py

@@ -16,6 +16,7 @@ def map(
     go_backwards=False,
     mode=None,
     name=None,
+    return_updates=True,
 ):
     """Construct a `Scan` `Op` that functions like `map`.
 
@@ -50,6 +51,7 @@ def map(
         go_backwards=go_backwards,
         mode=mode,
         name=name,
+        return_updates=return_updates,
     )
 
 
@@ -61,6 +63,7 @@ def reduce(
     go_backwards=False,
     mode=None,
     name=None,
+    return_updates=True,
 ):
     """Construct a `Scan` `Op` that functions like `reduce`.
 
@@ -97,14 +100,29 @@ def reduce(
         truncate_gradient=-1,
         mode=mode,
         name=name,
+        return_updates=return_updates,
     )
-    if isinstance(rval[0], list | tuple):
-        return [x[-1] for x in rval[0]], rval[1]
+    if return_updates:
+        if isinstance(rval[0], list | tuple):
+            return [x[-1] for x in rval[0]], rval[1]
+        else:
+            return rval[0][-1], rval[1]
     else:
-        return rval[0][-1], rval[1]
+        if isinstance(rval, list | tuple):
+            return [x[-1] for x in rval]
+        else:
+            return rval[-1]
 
 
-def foldl(fn, sequences, outputs_info, non_sequences=None, mode=None, name=None):
+def foldl(
+    fn,
+    sequences,
+    outputs_info,
+    non_sequences=None,
+    mode=None,
+    name=None,
+    return_updates=True,
+):
     """Construct a `Scan` `Op` that functions like Haskell's `foldl`.
 
     Parameters
@@ -135,10 +153,19 @@ def foldl(fn, sequences, outputs_info, non_sequences=None, mode=None, name=None)
         go_backwards=False,
         mode=mode,
         name=name,
+        return_updates=return_updates,
     )
 
 
-def foldr(fn, sequences, outputs_info, non_sequences=None, mode=None, name=None):
+def foldr(
+    fn,
+    sequences,
+    outputs_info,
+    non_sequences=None,
+    mode=None,
+    name=None,
+    return_updates=True,
+):
     """Construct a `Scan` `Op` that functions like Haskell's `foldr`.
 
     Parameters
@@ -169,4 +196,5 @@ def foldr(fn, sequences, outputs_info, non_sequences=None, mode=None, name=None)
         go_backwards=True,
         mode=mode,
         name=name,
+        return_updates=return_updates,
     )

pytensor/tensor/pad.py

@@ -314,11 +314,12 @@ def _wrap_pad(x: TensorVariable, pad_width: TensorVariable) -> TensorVariable:
 
 
 def _build_padding_one_direction(array, array_flipped, repeats, *, inner_func, axis):
-    [_, parts], _ = scan(
+    [_, parts] = scan(
         inner_func,
         non_sequences=[array, array_flipped],
         outputs_info=[0, None],
         n_steps=repeats,
+        return_updates=False,
     )
 
     parts = moveaxis(parts, 0, axis)

tests/scan/test_basic.py

@@ -27,7 +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.basic import Apply, equal_computations
+from pytensor.graph.basic import Apply, Variable, equal_computations
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import Op
 from pytensor.graph.replace import vectorize_graph
@@ -67,6 +67,7 @@
     vector,
 )
 from tests import unittest_tools as utt
+from tests.unittest_tools import assert_equal_computations
 
 
 if config.mode == "FAST_COMPILE":
@@ -4139,3 +4140,46 @@ def step(prev_x, prev_rng):
         xs_ref.append(rng_ref.normal(xs_ref[-1]))
     assert random_generator_type.values_eq(rng_ref, rng_final_eval)
     np.testing.assert_allclose(xs_eval, xs_ref[1:])
+
+
+@pytest.mark.filterwarnings("error")
+def test_return_updates_api_change():
+    err_msg = "return_updates=False but Scan produced updates"
+    warn_msg1 = "Updates functionality in Scan are deprecated"
+    warn_msg2 = "Pass `return_updates=False` to conform to the new API"
+
+    x = shared(np.array(0, dtype="float64"))
+
+    with pytest.warns(DeprecationWarning, match=warn_msg2):
+        with pytest.warns(DeprecationWarning, match=warn_msg1):
+            traced1, updates1 = scan(
+                lambda: {x: x + 1},
+                outputs_info=[],
+                n_steps=5,
+            )
+    assert traced1 is None
+    assert len(updates1) == 1 and x in updates1
+
+    with pytest.warns(DeprecationWarning, match=warn_msg2):
+        traced2, updates2 = scan(
+            lambda x: x + 1,
+            outputs_info=[x],
+            n_steps=5,
+        )
+    assert isinstance(traced2, Variable)
+    assert isinstance(updates2, dict) and not updates2
+
+    traced3 = scan(
+        lambda x: x + 1,
+        outputs_info=[x],
+        n_steps=5,
+        return_updates=False,
+    )
+    assert isinstance(traced3, Variable)
+
+    assert_equal_computations(list(updates1.values()), [traced2[-1]])
+    assert_equal_computations([traced2], [traced3])
+
+    with pytest.raises(ValueError, match=err_msg):
+        with pytest.warns(DeprecationWarning, match=warn_msg1):
+            scan(lambda: {x: x + 1}, outputs_info=[], n_steps=5, return_updates=False)