Support while loops in the JAX dispatcher

Author: rlouf

aesara/link/jax/dispatch/scan.py

@@ -36,13 +36,61 @@ def parse_outer_inputs(outer_inputs):
         return outer_in
 
     if op.info.as_while:
-        raise NotImplementedError("While loops are not supported in the JAX backend.")
+        # The inner function returns a boolean as the last value.
+        return make_jax_while_fn(scan_inner_fn, parse_outer_inputs, input_taps)
     else:
-        return make_jax_scan_fn(
-            scan_inner_fn,
-            parse_outer_inputs,
-            input_taps,
-        )
+        return make_jax_scan_fn(scan_inner_fn, parse_outer_inputs, input_taps)
+
+
+def make_jax_while_fn(
+    scan_inner_fn: Callable,
+    parse_outer_inputs: Callable[[TensorVariable], Dict[str, List[TensorVariable]]],
+    input_taps: Dict,
+):
+    """Create a `jax.lax.while_loop` function to perform `Scan` computations when it
+    is used as while loop.
+
+    `jax.lax.while_loop` iterates by passing a value `carry` to a `body_fun` that
+    must return a value of the same type (Pytree structure, shape and dtype of
+    the leaves). Before calling `body_fn`, it calls `cond_fn` which takes the
+    current value and returns a boolean that indicates whether to keep iterating
+    or not.
+
+    The JAX `while_loop` needs to perform the following operations:
+
+    1. Extract the inner-inputs;
+    2. Build the initial carry value;
+    3. Inside the loop:
+        1. `carry` -> inner-inputs;
+        2. inner-outputs -> `carry`
+    4. Post-process the `carry`  storage and return outputs
+    """
+
+    def build_while_carry(outer_in):
+        return outer_in
+
+    def while_loop(*outer_inputs):
+
+        outer_in = parse_outer_inputs(outer_inputs)
+        init_carry = build_while_carry(outer_in)
+
+        def cond_fn(carry):
+            # The inner-function of `Scan` returns a boolean as the last
+            # value. This needs to be included in `carry`.
+            # TODO: Will it return `False` if the number of steps is exceeded?
+            return carry["do_continue"]
+
+        def body_fn(carry):
+            return carry
+
+        carry = jax.lax.while_loop(body_fn, cond_fn, init_carry)
+
+        # Here post-process the result
+        outer_outputs = carry
+
+        return outer_outputs
+
+    return while_loop
 
 
 def make_jax_scan_fn(
@@ -58,7 +106,8 @@ def make_jax_scan_fn(
     stacked to the previous outputs. We use this to our advantage to build
     `Scan` outputs without having to post-process the storage arrays.
 
-    The JAX scan function needs to perform the following operations:
+    The JAX `scan` function needs to perform the following operations:
+
     1. Extract the inner-inputs;
     2. Build the initial `carry` and `sequence` values;
     3. Inside the loop:
@@ -265,11 +314,11 @@ def body_fn(carry, x):
         )
 
         shared_output = tuple(last_carry["shared"])
-        results = results + shared_output
+        outer_outputs = results + shared_output
 
-        if len(results) == 1:
-            return results[0]
+        if len(outer_outputs) == 1:
+            return outer_outputs[0]
 
-        return results
+        return outer_outputs
 
     return scan

tests/link/jax/test_scan.py

@@ -11,6 +11,7 @@
 from aesara.link.jax.linker import JAXLinker
 from aesara.scan.basic import scan
 from aesara.scan.op import Scan
+from aesara.scan.utils import until
 from aesara.tensor.math import gammaln, log
 from aesara.tensor.random.utils import RandomStream
 from aesara.tensor.type import ivector, lscalar, scalar
@@ -24,6 +25,15 @@
 jax_mode = Mode(JAXLinker(), opts)
 
 
+def test_while():
+    res, updates = scan(
+        fn=lambda a_tm1: (a_tm1 + 1, until(a_tm1 > 2)),
+        outputs_info=[{"initial": at.as_tensor(1, dtype=np.int64), "taps": [-1]}],
+        n_steps=3,
+    )
+    jax_fn = function((), res, updates=updates, mode="JAX")
+
+
 def test_sit_sot():
     a_at = at.scalar("a", dtype="floatX")