Return only the last value of the shared variables

Author: rlouf

aesara/link/jax/dispatch/scan.py

@@ -177,13 +177,50 @@ def scan_new_carry(carry, inner_outputs):
 
             return new_carry
 
+        def scan_new_outputs(inner_outputs):
+            """Create a new outer-output value from the inner-output value.
+
+            Outer-outputs are ordered as follows:
+            - mit-mot-outputs
+            - mit-sot-outputs
+            - sit-sot-outputs
+            - nit-sots
+            - shared-outputs
+
+            The shared output corresponds to the last value found in the last
+            carry value returned by `jax.lax.scan`. It is thus not returned in
+            the body function.
+
+            """
+            outer_outputs = []
+            if "mit_sot" in inner_output_idx:
+                outer_outputs.append(
+                    [inner_outputs[idx] for idx in inner_output_idx["mit_sot"]]
+                )
+            if "sit_sot" in inner_output_idx:
+                outer_outputs.append(
+                    [inner_outputs[idx] for idx in inner_output_idx["sit_sot"]]
+                )
+            if "nit_sot" in inner_output_idx:
+                outer_outputs.append(
+                    [inner_outputs[idx] for idx in inner_output_idx["nit_sot"]]
+                )
+
+            return tuple(sum(outer_outputs, []))
+
         def body_fn(carry, x):
             inner_in_args = scan_inner_in_args(carry, x)
             inner_outputs = scan_inner_func(*inner_in_args)
             new_carry = scan_new_carry(carry, inner_outputs)
-            return new_carry, inner_outputs
+            outer_outputs = scan_new_outputs(inner_outputs)
+            return new_carry, outer_outputs
+
+        last_carry, results = jax.lax.scan(
+            body_fn, init_carry, sequences, length=n_steps
+        )
 
-        _, results = jax.lax.scan(body_fn, init_carry, sequences, length=n_steps)
+        shared_output = tuple(last_carry["shared"])
+        results = results + shared_output
 
         if len(results) == 1:
             return results[0]