Repository: google/flax

Issue:
Support `unrolled` steps in `jax_utils.scan_in_dims`
Motivated by [jax#3094](https://github.com/google/jax/issues/3094), [jax#3738](https://github.com/google/jax/pull/3738) and [jax#3076](https://github.com/google/jax/pull/3076), `jax.lax.scan` currently supports specifying the number of scan iterations to unroll into a single iteration of the loop using the argument `unrolls`.

It would be nice to be able to control this from `jax_utils.scan_in_dims`.


Discussion:

Resolving patch:
diff --git a/flax/jax_utils.py b/flax/jax_utils.py
--- a/flax/jax_utils.py
+++ b/flax/jax_utils.py
@@ -159,7 +159,7 @@ def enqueue(n):  # Enqueues *up to* `n` elements from the iterator.
     enqueue(1)
 
 
-def _scan_nd(body_fn, init, xs, n=1):
+def _scan_nd(body_fn, init, xs, n=1, unroll=(1,)):
   """Utility for performing an n-dimensional `lax.scan`.
 
   The n-d scan is simply recursive call of 1-d scan.
@@ -172,11 +172,11 @@ def _scan_nd(body_fn, init, xs, n=1):
     A tuple of the final carry and the values returned by the body.
   """
   if n == 1:
-    return lax.scan(body_fn, init, xs)
+    return lax.scan(body_fn, init, xs, unroll=unroll[0])
   else:
     def scan_body(c, x):
-      return _scan_nd(body_fn, c, x, n=n-1)
-    return lax.scan(scan_body, init, xs)
+      return _scan_nd(body_fn, c, x, n=n-1, unroll=unroll[1:])
+    return lax.scan(scan_body, init, xs, unroll=unroll[0])
 
 
 def _invert_perm(perm):
@@ -186,22 +186,38 @@ def _invert_perm(perm):
   return tuple(perm_inv)
 
 
-def scan_in_dim(body_fn, init, xs, axis=(0,), keepdims=False):
+def scan_in_dim(body_fn, init, xs, axis=(0,), unroll=(1,), keepdims=False):
   """utility for doing a scan along arbitrary dimensions.
 
-  see `lax.scan` for details on how the scan operation works.
+  See `lax.scan` for details on how the scan operation works.
+
+  Note on `unroll`: This argument gets left padded with ones to match the size
+  of `axis`. Doing so allows unrolls to performed from the innermost loop first.
+  For example, `scan_in_dim(..., axis=(1, 2, 3), unroll=5)` is equivalent to
+  `scan_in_dim(..., axis=(1, 2, 3), unroll=(1, 1, 5))`.
+
   Args:
     body_fn: the body of the loop of type (c, x) -> (c, y).
     init: initial value for the carry.
     xs: a pytree of tensors to scan over.
     axis: the axis to scan over.
     keepdims: keep the dimensions that are scanned over.
+    unroll: an optional positive integer, or tuple of positive integers
+      showing how many iterations of the loop to be unroll into a single
+      iteration for each axis.
   Returns:
     A tuple of the final carry and the values returned by the body.
   """
   if not isinstance(axis, Iterable):
     axis = (axis,)
 
+  if not isinstance(unroll, Iterable):
+    unroll = (unroll,)
+
+  # Pad unroll with ones so we start unrolling from the innermost loop
+  len_diff = len(axis) - len(unroll)
+  unroll = (1,) * len_diff + unroll
+
   def transpose_in(x):
     perm = axis + tuple(np.delete(np.arange(x.ndim), axis))
     return x.transpose(perm)
@@ -220,6 +236,6 @@ def body_wrapper(c, xs):
     return c, ys
 
   xs = jax.tree_map(transpose_in, xs)
-  c, ys = _scan_nd(body_wrapper, init, xs, n=len(axis))
+  c, ys = _scan_nd(body_wrapper, init, xs, n=len(axis), unroll=unroll)
   ys = jax.tree_map(transpose_out, ys)
   return c, ys


Test patch: