Manage indexing when there are several xit-sots

Author: rlouf

aesara/link/jax/dispatch/scan.py

@@ -66,11 +66,7 @@ def scan(*outer_inputs):
         n_steps = outer_in["n_steps"]
         sequences = outer_in["sequences"]
         non_sequences = outer_in["non_sequences"]
-        init_carry = {
-            name: outer_in[name]
-            for name in ["mit_sot", "sit_sot", "shared"]
-            if len(outer_in[name]) > 0
-        }
+        init_carry = {name: outer_in[name] for name in ["mit_sot", "sit_sot", "shared"]}
         init_carry["step"] = 0
 
         # Map to retrieve the inner-outputs
@@ -100,16 +96,18 @@ def scan_inner_in_args(carry, x):
             inner_in_mit_sot = sum(
                 [
                     convert(step, carry_element)
-                    for convert in mit_sot_from_carry
-                    for carry_element in carry["mit_sot"]
+                    for convert, carry_element in zip(
+                        mit_sot_from_carry, carry["mit_sot"]
+                    )
                 ],
                 [],
             )
             inner_in_sit_sot = sum(
                 [
                     convert(step, carry_element)
-                    for convert in sit_sot_from_carry
-                    for carry_element in carry["sit_sot"]
+                    for convert, carry_element in zip(
+                        sit_sot_from_carry, carry["sit_sot"]
+                    )
                 ],
                 [],
             )
@@ -138,8 +136,8 @@ def scan_new_carry(carry, inner_outputs):
             [+ while-condition]
 
             """
-            new_carry = {}
             step = carry["step"]
+            new_carry = {"mit_sot": [], "sit_sot": [], "shared": []}
 
             if "shared" in inner_output_idx:
                 shared_inner_outputs = [
@@ -154,9 +152,8 @@ def scan_new_carry(carry, inner_outputs):
                 new_carry["mit_sot"] = sum(
                     [
                         convert(step, carry_element, inner_outputs_element)
-                        for convert in mit_sot_to_carry
-                        for (carry_element, inner_outputs_element) in zip(
-                            carry["mit_sot"], mit_sot_inner_outputs
+                        for (convert, carry_element, inner_outputs_element) in zip(
+                            mit_sot_to_carry, carry["mit_sot"], mit_sot_inner_outputs
                         )
                     ],
                     [],
@@ -169,9 +166,8 @@ def scan_new_carry(carry, inner_outputs):
                 new_carry["sit_sot"] = sum(
                     [
                         convert(step, 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
+                        for (convert, carry_element, inner_outputs_element) in zip(
+                            sit_sot_to_carry, carry["sit_sot"], sit_sot_inner_outputs
                         )
                     ],
                     [],
@@ -189,6 +185,9 @@ def body_fn(carry, x):
 
         _, results = jax.lax.scan(body_fn, init_carry, sequences, length=n_steps)
 
-        return results[0]
+        if len(results) == 1:
+            return results[0]
+
+        return results
 
     return scan

tests/link/jax/test_scan.py

@@ -69,6 +69,22 @@ def test_mit_sot():
     assert np.allclose(fn(), jax_fn())
 
 
+def test_mit_sot_2():
+    res, updates = scan(
+        fn=lambda a_tm1, b_tm1: (2 * a_tm1, 2 * b_tm1),
+        outputs_info=[
+            {"initial": at.as_tensor(1.0, dtype="floatX"), "taps": [-1]},
+            {"initial": at.as_tensor(0.5, dtype="floatX"), "taps": [-1]},
+        ],
+        n_steps=10,
+    )
+    jax_fn = function((), res, updates=updates, mode="JAX")
+    fn = function((), res, updates=updates)
+    print(jax_fn())
+    print(fn())
+    assert np.allclose(fn(), jax_fn())
+
+
 @pytest.mark.parametrize(
     "fn, sequences, outputs_info, non_sequences, n_steps, input_vals, output_vals, op_check",
     [
@@ -126,6 +142,44 @@ def test_mit_sot():
         #     None,
         #     lambda op: op.info.n_sit_sot > 0,
         # ),
+        # # nit-sot, shared input/output
+        # (
+        #     lambda: RandomStream(seed=1930, rng_ctor=np.random.RandomState).normal(
+        #         0, 1, name="a"
+        #     ),
+        #     [],
+        #     [{}],
+        #     [],
+        #     3,
+        #     [],
+        #     [np.array([-1.63408257, 0.18046406, 2.43265803])],
+        #     lambda op: op.info.n_shared_outs > 0,
+        # ),
+        # # mit-sot (that's also a type of sit-sot)
+        # (
+        #     lambda a_tm1: 2 * a_tm1,
+        #     [],
+        #     [{"initial": at.as_tensor([0.0, 1.0], dtype="floatX"), "taps": [-2]}],
+        #     [],
+        #     6,
+        #     [],
+        #     None,
+        #     lambda op: op.info.n_mit_sot > 0,
+        # ),
+        # # mit-sot
+        # (
+        #     lambda a_tm1, b_tm1: (2 * a_tm1, 2 * b_tm1),
+        #     [],
+        #     [
+        #         {"initial": at.as_tensor(1.0, dtype="floatX"), "taps": [-1]},
+        #         {"initial": at.as_tensor(0.3, dtype="floatX"), "taps": [-1]},
+        #     ],
+        #     [],
+        #     10,
+        #     [],
+        #     None,
+        #     lambda op: op.info.n_mit_sot > 0,
+        # ),
     ],
 )
 def test_xit_xot_types(