Use jax.lax.dynamic_slice in the Subtensor JAX dispatcher

Author: rlouf

aesara/link/jax/dispatch/subtensor.py

@@ -1,5 +1,8 @@
+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,
@@ -12,21 +15,55 @@
 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, **kwargs):
+def jax_funcify_Subtensor(op, node, **kwargs):
 
     idx_list = getattr(op, "idx_list", None)
 
-    def subtensor(x, *ilists):
+    # 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]
 
-        return x.__getitem__(indices)
+        if isinstance(indices, slice):
+            return jax_dynamic_slice(x, indices.start, indices.stop)
+        else:
+            index = indices
+            return x.__getitem__(index)
 
     return subtensor