Start deprecating updates API

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

@@ -163,6 +163,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,
@@ -177,6 +197,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.
 
@@ -873,6 +894,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:
@@ -895,7 +921,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
@@ -914,6 +940,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
@@ -1206,12 +1234,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
@@ -1244,4 +1274,4 @@ def remove_dimensions(outs, offsets=None):
     elif len(scan_out_list) == 0:
         scan_out_list = None
 
-    return scan_out_list, update_map
+    return _manage_output_api_change(scan_out_list, update_map, 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/link/jax/test_scan.py

@@ -23,10 +23,11 @@
 @pytest.mark.parametrize("view", [None, (-1,), slice(-2, None, None)])
 def test_scan_sit_sot(view):
     x0 = pt.scalar("x0", dtype="float64")
-    xs, _ = scan(
+    xs = scan(
         lambda xtm1: xtm1 + 1,
         outputs_info=[x0],
         n_steps=10,
+        return_updates=False,
     )
     if view:
         xs = xs[view]

tests/link/numba/test_scan.py

@@ -343,8 +343,12 @@ def power_step(prior_result, x):
 
 def test_grad_sitsot():
     def get_sum_of_grad(inp):
-        scan_outputs, _updates = scan(
-            fn=lambda x: x * 2, outputs_info=[inp], n_steps=5, mode="NUMBA"
+        scan_outputs = scan(
+            fn=lambda x: x * 2,
+            outputs_info=[inp],
+            n_steps=5,
+            mode="NUMBA",
+            return_updates=False,
         )
         return grad(scan_outputs.sum(), inp).sum()