Create a dummy initial carry value when outputs_info is empty

Author: rlouf

aesara/link/jax/dispatch/scan.py

@@ -15,7 +15,7 @@ def scan(*outer_inputs):
         outer_in_mit_mot = list(op.outer_mitmot(outer_inputs))
         outer_in_mit_sot = list(op.outer_mitsot(outer_inputs))
         outer_in_sit_sot = list(op.outer_sitsot(outer_inputs))
-        outer_in_nit_sot = list(op.outer_nitsot(outer_inputs))
+        # outer_in_nit_sot = list(op.outer_nitsot(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):
@@ -29,10 +29,18 @@ def scan(*outer_inputs):
         if len(outer_in_non_seqs):
             raise NotImplementedError("non sequence are not supported")
 
-        init_carry = outer_in_nit_sot
+        # If `output_infos` is empty we need to create an empty initial carry
+        # value with the output's shape and dtype
+        is_outputs_info_empty = len(outer_in_sit_sot) == 0
+        if is_outputs_info_empty:
+            dtype = node.outputs[0].type.dtype
+            shape = tuple(s for s in node.outputs[0].type.shape if s is not None)
+            outer_in_sit_sot = [jax.numpy.empty(dtype=dtype, shape=shape)]
+
+        init_carry = outer_in_sit_sot
         sequences = outer_in_seqs
 
-        def scan_inner_in_args(carry, x, is_dummy_sit_sot=True):
+        def scan_inner_in_args(carry, x):
             """Create an inner-input expression.
 
             Inner-inputs are ordered as follows:
@@ -45,10 +53,12 @@ def scan_inner_in_args(carry, x, is_dummy_sit_sot=True):
             """
 
             inner_in_seqs = x
-            if is_dummy_sit_sot:
+
+            if is_outputs_info_empty:
                 inner_in_sit_sot = []
             else:
                 inner_in_sit_sot = carry
+
             return sum([inner_in_seqs, inner_in_sit_sot], [])
 
         def scan_new_carry(inner_outputs):
@@ -58,7 +68,7 @@ def scan_new_carry(inner_outputs):
             - mit-mot-outputs
             - mit-sot-outputs
             - sit-sot-outputs
-            - nit-sots
+            - nit-sots (no carry)
             - shared-outputs
             [+ while-condition]
 

aesara/link/jax/dispatch/subtensor.py

@@ -1,8 +1,5 @@
-import jax
-import jax.numpy as jnp
 
 from aesara.link.jax.dispatch.basic import jax_funcify
-from aesara.scalar.basic import ScalarConstant
 from aesara.tensor.subtensor import (
     AdvancedIncSubtensor,
     AdvancedIncSubtensor1,
@@ -15,55 +12,22 @@
 from aesara.tensor.type_other import MakeSlice
 
 
-def jax_dynamic_slice(
-    array_to_slice: jnp.DeviceArray, start: jnp.DeviceArray, stop: jnp.DeviceArray
-) -> jnp.DeviceArray:
-    """Slice a JAX array with traced indices with jax.lax.dynamic_slice.
-
-    TODO: This currently assumes that arrays are 1-dimensional.
-    """
-    if stop is None:
-        slice_size = array_to_slice.shape[0] - start
-    if start is None:
-        start = 0
-        slice_size = stop
-    return jax.lax.dynamic_slice(array_to_slice, (start,), (slice_size,))
-
-
 @jax_funcify.register(Subtensor)
 @jax_funcify.register(AdvancedSubtensor)
 @jax_funcify.register(AdvancedSubtensor1)
 def jax_funcify_Subtensor(op, node, **kwargs):
 
     idx_list = getattr(op, "idx_list", None)
 
-    # JAX does not provide an easy way to slice an array dynamically with
-    # a step value, so we raise an exception.
-    #
-    # However, Aesara rewrites sometimes introduce a constant step with value
-    # `1` in which case it is safe to convert the Op, ignoring the step value.
-    if len(node.inputs) == 4:
-        if not isinstance(node.inputs[-1], ScalarConstant):
-            raise NotImplementedError(
-                "Dynamic slicing operations with a step value are not supported in JAX"
-            )
-        else:
-            if node.inputs[-1].value != 1:
-                raise NotImplementedError(
-                    "Dynamic slicing operations with a step value are not supported in JAX"
-                )
 
     def subtensor(x, *ilists):
+
         indices = indices_from_subtensor(ilists, idx_list)
 
         if len(indices) == 1:
             indices = indices[0]
 
-        if isinstance(indices, slice):
-            return jax_dynamic_slice(x, indices.start, indices.stop)
-        else:
-            index = indices
-            return x.__getitem__(index)
+        return x.__getitem__(indices)
 
     return subtensor
 

tests/link/jax/test_scan.py

@@ -37,6 +37,17 @@
             None,
             lambda op: op.info.n_seqs > 0,
         ),
+        # nit-sot
+        (
+            lambda: at.as_tensor(2.0),
+            [],
+            [{}],
+            [],
+            3,
+            [],
+            None,
+            lambda op: op.info.n_nit_sot > 0,
+        ),
     ],
 )
 def test_xit_xot_types(