Manage indexing of carried arrays for mit-sots

Author: rlouf

aesara/link/jax/dispatch/scan.py

@@ -1,3 +1,5 @@
+from collections import defaultdict
+
 import jax
 
 from aesara.link.jax.dispatch.basic import jax_funcify
@@ -10,7 +12,8 @@ def jax_funcify_Scan(op, node, **kwargs):
     if op.info.as_while:
         raise NotImplementedError("While loops are not supported in the JAX backend.")
 
-    input_taps = op.info.sit_sot_in_slices
+    sit_sot_input_taps = op.info.sit_sot_in_slices
+    mit_sot_input_taps = op.info.mit_sot_in_slices
 
     # Construct `scan_inner_func`'s arguments from the carry value and sequence
     # element passed to `body_fn`.
@@ -19,69 +22,65 @@ def jax_funcify_Scan(op, node, **kwargs):
     # inputs taps.
     def index_carry_arrays(input_taps):
         """Fetch the inner inputs from the values stored in the carry array"""
+        # TODO: Check and refactor this
         storage_size = -min(input_taps)
         offsets = [storage_size + tap for tap in input_taps]
 
-        def to_inner_inputs(carry):
-            return [carry[offset] for offset in offsets]
+        def to_inner_inputs(step, carry):
+            return [carry[step + offset] for offset in offsets]
 
         return to_inner_inputs
 
-    sit_sot_from_carry = [index_carry_arrays(tap) for tap in input_taps]
+    sit_sot_from_carry = [index_carry_arrays(tap) for tap in sit_sot_input_taps]
+    mit_sot_from_carry = [index_carry_arrays(tap) for tap in mit_sot_input_taps]
 
     # Construct the new carry values from the outputs of `scan_inner_func`
     def inner_outputs_to_carry(input_taps):
         """Create the new carry array from the inner output"""
-        storage_size = -min(input_taps)
-        offsets = [storage_size + tap for tap in input_taps]
-
-        def to_new_carry(carry, inner_outputs):
-            return [carry.at[offset].set(inner_outputs) for offset in offsets]
+        # TODO: Check and refactor this
+        def to_new_carry(step, carry, inner_outputs):
+            return [carry.at[step - tap].set(inner_outputs) for tap in input_taps]
 
         return to_new_carry
 
-    sit_sot_to_carry = [inner_outputs_to_carry(tap) for tap in input_taps]
+    sit_sot_to_carry = [inner_outputs_to_carry(tap) for tap in sit_sot_input_taps]
+    mit_sot_to_carry = [inner_outputs_to_carry(tap) for tap in mit_sot_input_taps]
 
     def scan(*outer_inputs):
 
-        n_steps = outer_inputs[0]
-        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):
+        # Inputs to `aesara.scan`
+        outer_in = {
+            "n_steps": outer_inputs[0],
+            "sequences": list(op.outer_seqs(outer_inputs)),
+            "mit_mot": list(op.outer_mitmot(outer_inputs)),
+            "mit_sot": list(op.outer_mitsot(outer_inputs)),
+            "nit_sot": list(op.outer_nitsot(outer_inputs)),
+            "sit_sot": list(op.outer_sitsot(outer_inputs)),
+            "shared": list(op.outer_shared(outer_inputs)),
+            "non_sequences": list(op.outer_non_seqs(outer_inputs)),
+        }
+        if len(outer_in["mit_mot"]) > 0:
             raise NotImplementedError("mit-mot not supported")
-        if len(outer_in_mit_sot):
-            raise NotImplementedError("mit-sot not supported")
-
-        # 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
 
+        # Inputs to `jax.lax.scan`
+        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["step"] = 0
+
+        # Map to retrieve the inner-outputs
         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)
+        for name in ["mit_sot", "sit_sot", "nit_sot", "shared"]:
+            num_outputs = len(outer_in[name])
+            for i in range(num_outputs):
+                inner_output_idx[name].append(offset + i)
+            offset += num_outputs
 
         def scan_inner_in_args(carry, x):
             """Get inner-inputs from the arguments passed to the `jax.lax.scan` body function.
@@ -95,10 +94,20 @@ def scan_inner_in_args(carry, x):
             - non-sequences
 
             """
+            step = carry["step"]
+
             inner_in_seqs = x
+            inner_in_mit_sot = sum(
+                [
+                    convert(step, carry_element)
+                    for convert in mit_sot_from_carry
+                    for carry_element in carry["mit_sot"]
+                ],
+                [],
+            )
             inner_in_sit_sot = sum(
                 [
-                    convert(carry_element)
+                    convert(step, carry_element)
                     for convert in sit_sot_from_carry
                     for carry_element in carry["sit_sot"]
                 ],
@@ -107,7 +116,14 @@ def scan_inner_in_args(carry, x):
             inner_in_shared = carry.get("shared", [])
 
             return sum(
-                [inner_in_seqs, inner_in_sit_sot, inner_in_shared, non_sequences], []
+                [
+                    inner_in_seqs,
+                    inner_in_mit_sot,
+                    inner_in_sit_sot,
+                    inner_in_shared,
+                    non_sequences,
+                ],
+                [],
             )
 
         def scan_new_carry(carry, inner_outputs):
@@ -123,20 +139,36 @@ def scan_new_carry(carry, inner_outputs):
 
             """
             new_carry = {}
+            step = carry["step"]
 
             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 "mit_sot" in inner_output_idx:
+                mit_sot_inner_outputs = [
+                    inner_outputs[idx] for idx in inner_output_idx["mit_sot"]
+                ]
+                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
+                        )
+                    ],
+                    [],
+                )
+
             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)
+                        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
@@ -145,6 +177,8 @@ def scan_new_carry(carry, inner_outputs):
                     [],
                 )
 
+            new_carry["step"] = carry["step"] + 1
+
             return new_carry
 
         def body_fn(carry, x):

tests/link/jax/test_scan.py

@@ -38,6 +38,9 @@ def test_sit_sot():
     assert np.allclose(fn(1.0), jax_fn(1.0))
 
 
+@pytest.mark.xfail(
+    reason="Returns correct results but raises exception due to stucture of shared variable."
+)
 def test_nit_sot_shared():
     res, updates = scan(
         fn=lambda: RandomStream(seed=1930, rng_ctor=np.random.RandomState).normal(
@@ -52,6 +55,20 @@ def test_nit_sot_shared():
     print(fn())
 
 
+def test_mit_sot():
+    res, updates = scan(
+        fn=lambda a_tm1: 2 * a_tm1,
+        outputs_info=[
+            {"initial": at.as_tensor([0.0, 1.0], dtype="floatX"), "taps": [-2]}
+        ],
+        n_steps=6,
+    )
+
+    jax_fn = function((), res, updates=updates, mode="JAX")
+    fn = function((), res, updates=updates)
+    assert np.allclose(fn(), jax_fn())
+
+
 @pytest.mark.parametrize(
     "fn, sequences, outputs_info, non_sequences, n_steps, input_vals, output_vals, op_check",
     [
@@ -96,7 +113,6 @@ def test_nit_sot_shared():
         #     [],
         #     3,
         #     [],
-        #     None,
         #     lambda op: op.info.n_sit_sot > 0,
         # ),
         # # sit-sot, while