Support shared inputs in JAX's Scan dispatch

Author: rlouf

aesara/link/jax/dispatch/scan.py

@@ -48,19 +48,40 @@ def scan(*outer_inputs):
         outer_in_seqs = list(op.outer_seqs(outer_inputs))
         outer_in_mit_mot = list(op.outer_mitmot(outer_inputs))
         outer_in_mit_sot = list(op.outer_mitsot(outer_inputs))
+        outer_in_nit_sot = list(op.outer_nitsot(outer_inputs))
         outer_in_sit_sot = list(op.outer_sitsot(outer_inputs))
         outer_in_shared = list(op.outer_shared(outer_inputs))
         outer_in_non_seqs = list(op.outer_non_seqs(outer_inputs))
         if len(outer_in_mit_mot):
             raise NotImplementedError("mit-mot not supported")
         if len(outer_in_mit_sot):
             raise NotImplementedError("mit-sot not supported")
-        if len(outer_in_shared):
-            raise NotImplementedError("shared variables not supported")
 
-        init_carry = outer_in_sit_sot
+        # These are the outer-inputs
         sequences = outer_in_seqs
         non_sequences = outer_in_non_seqs
+        init_carry = {}
+        for name, outputs in [
+            ("sit_sot", outer_in_sit_sot),
+            ("shared", outer_in_shared),
+        ]:
+            if len(outputs) > 0:
+                init_carry[name] = outputs
+
+        # We keep track of the kind of inner_outputs and their number
+        from collections import defaultdict
+
+        offset = 0
+        inner_output_idx = defaultdict(list)
+        for name, outputs in [
+            ("sit_sot", outer_in_sit_sot),
+            ("nit_sot", outer_in_nit_sot),
+            ("shared", outer_in_shared),
+        ]:
+            if len(outputs) > 0:
+                for i in range(len(outputs)):
+                    inner_output_idx[name].append(offset + i)
+                offset += len(outputs)
 
         def scan_inner_in_args(carry, x):
             """Get inner-inputs from the arguments passed to the `jax.lax.scan` body function.
@@ -72,18 +93,22 @@ def scan_inner_in_args(carry, x):
             - sit-sot inputs
             - shared-inputs
             - non-sequences
+
             """
             inner_in_seqs = x
             inner_in_sit_sot = sum(
                 [
                     convert(carry_element)
                     for convert in sit_sot_from_carry
-                    for carry_element in carry
+                    for carry_element in carry["sit_sot"]
                 ],
                 [],
             )
+            inner_in_shared = carry.get("shared", [])
 
-            return sum([inner_in_seqs, inner_in_sit_sot, non_sequences], [])
+            return sum(
+                [inner_in_seqs, inner_in_sit_sot, inner_in_shared, non_sequences], []
+            )
 
         def scan_new_carry(carry, inner_outputs):
             """Create a new carry value from the inner-outputs.
@@ -97,27 +122,39 @@ def scan_new_carry(carry, inner_outputs):
             [+ while-condition]
 
             """
-            carry = sum(
-                [
-                    convert(carry_element, inner_outputs_element)
-                    for convert in sit_sot_to_carry
-                    for (carry_element, inner_outputs_element) in zip(
-                        carry, inner_outputs
-                    )
-                ],
-                [],
-            )
-            return carry
+            new_carry = {}
+
+            if "shared" in inner_output_idx:
+                shared_inner_outputs = [
+                    inner_outputs[idx] for idx in inner_output_idx["shared"]
+                ]
+                new_carry["shared"] = shared_inner_outputs
+
+            if "sit_sot" in inner_output_idx:
+                sit_sot_inner_outputs = [
+                    inner_outputs[idx] for idx in inner_output_idx["sit_sot"]
+                ]
+                new_carry["sit_sot"] = sum(
+                    [
+                        convert(carry_element, inner_outputs_element)
+                        for convert in sit_sot_to_carry
+                        for (carry_element, inner_outputs_element) in zip(
+                            carry["sit_sot"], sit_sot_inner_outputs
+                        )
+                    ],
+                    [],
+                )
+
+            return new_carry
 
         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
+            return new_carry, inner_outputs
 
-        print(init_carry)
         _, results = jax.lax.scan(body_fn, init_carry, sequences, length=n_steps)
 
-        return results
+        return results[0]
 
     return scan

tests/link/jax/test_scan.py

@@ -12,6 +12,7 @@
 from aesara.scan.basic import scan
 from aesara.scan.op import Scan
 from aesara.tensor.math import gammaln, log
+from aesara.tensor.random.utils import RandomStream
 from aesara.tensor.type import ivector, lscalar, scalar
 from tests.link.jax.test_basic import compare_jax_and_py
 
@@ -37,6 +38,20 @@ def test_sit_sot():
     assert np.allclose(fn(1.0), jax_fn(1.0))
 
 
+def test_nit_sot_shared():
+    res, updates = scan(
+        fn=lambda: RandomStream(seed=1930, rng_ctor=np.random.RandomState).normal(
+            0, 1, name="a"
+        ),
+        n_steps=3,
+    )
+
+    jax_fn = function((), res, updates=updates, mode="JAX")
+    print(jax_fn())
+    fn = function((), res, updates=updates)
+    print(fn())
+
+
 @pytest.mark.parametrize(
     "fn, sequences, outputs_info, non_sequences, n_steps, input_vals, output_vals, op_check",
     [
@@ -73,6 +88,28 @@ def test_sit_sot():
             None,
             lambda op: op.info.n_nit_sot > 0 and op.info.n_non_seqs > 0,
         ),
+        # # sit-sot
+        # (
+        #     lambda a_tm1: 2 * a_tm1,
+        #     [],
+        #     [{"initial": at.as_tensor(0.0, dtype="floatX"), "taps": [-1]}],
+        #     [],
+        #     3,
+        #     [],
+        #     None,
+        #     lambda op: op.info.n_sit_sot > 0,
+        # ),
+        # # sit-sot, while
+        # (
+        #     lambda a_tm1: (a_tm1 + 1, until(a_tm1 > 2)),
+        #     [],
+        #     [{"initial": at.as_tensor(1, dtype=np.int64), "taps": [-1]}],
+        #     [],
+        #     3,
+        #     [],
+        #     None,
+        #     lambda op: op.info.n_sit_sot > 0,
+        # ),
     ],
 )
 def test_xit_xot_types(