Exposing the experimental _split_transpose JAX scan parameter in Flax.

flax/core/axes_scan.py

@@ -40,6 +40,7 @@ def scan(
   length: Optional[int] = None,
   reverse: bool = False,
   unroll: int = 1,
+  _split_transpose: bool = False
 ):
   """A wrapper around `jax.lax.scan` with in_axes/out_axes api.
 
@@ -74,6 +75,8 @@ def body_fn(b, c, x):
     reverse: scan in reverse order from end to start.
     unroll: how many scan iterations to unroll within a single
       iteration of a loop (default: 1).
+    _split_transpose: An experimental feature to split the transpose of scan
+       into a scan and a map, backed by an experimental Jax lax.scan() feature.
   Returns:
      the function that performs the scan of the form:
      (broadcast_in, carry_in, *args) -> (broadcast_out, carry_out, scan_out).
@@ -158,9 +161,15 @@ def body_fn(c, xs, init_mode=False):
       out_tree(), out_flat
     )
 
-    c, ys = lax.scan(
-      body_fn, init, xs, length=length, reverse=reverse, unroll=unroll
-    )
+    if jax.version.__version_info__ > (0, 4, 25):
+      c, ys = lax.scan(
+        body_fn, init, xs, length=length, reverse=reverse, unroll=unroll,
+        _split_transpose=_split_transpose
+      )
+    else:
+      c, ys = lax.scan(
+        body_fn, init, xs, length=length, reverse=reverse, unroll=unroll
+      )
     ys = jax.tree_util.tree_map(transpose_from_front, out_axes, ys)
     ys = jax.tree_util.tree_map(
       lambda ax, const, y: (const if ax is broadcast else y),

flax/core/lift.py

@@ -872,6 +872,7 @@ def scan(
   length: Optional[int] = None,
   reverse: bool = False,
   unroll: int = 1,
+  _split_transpose: bool = False,
   data_transform: Optional[Callable[..., Any]] = None,
   metadata_params: Dict[Any, Any] = {},
 ) -> Callable[..., Any]:
@@ -935,6 +936,8 @@ def body_fn(scope, c, x):
     reverse: If true, scan from end to start in reverse order.
     unroll: how many scan iterations to unroll within a single
       iteration of a loop (default: 1).
+    _split_transpose: An experimental feature to split the transpose of a scan
+       into a scan and a map, backed by an experimental Jax lax.scan() feature.
     data_transform: optional function to transform raw variable and rng groups,
       intended for inline SPMD annotations.
     metadata_params: arguments dict passed to AxisMetadata instances in the
@@ -993,6 +996,7 @@ def find_length(axis, x):
         length=length,
         reverse=reverse,
         unroll=unroll,
+        _split_transpose=_split_transpose
     )
     def scanned(broadcast_vars, carry, scan_variable_groups, rng_groups, args):
       carry_vars, c = carry

flax/linen/transforms.py

@@ -1143,6 +1143,7 @@ def scan(
   data_transform: Optional[Callable[..., Any]] = None,
   metadata_params: Mapping[Any, Any] = {},
   methods=None,
+  _split_transpose: bool = False,
 ) -> Target:
   """A lifted version of ``jax.lax.scan``.
 
@@ -1280,6 +1281,8 @@ def scan(
     metadata_params: arguments dict passed to AxisMetadata instances in the
       variable tree.
     methods: If ``target`` is a ``Module``, the methods of ``Module`` to scan over.
+    _split_transpose: An experimental feature to split the transpose of a scan
+       into a scan and a map, backed by an experimental Jax lax.scan() feature.
 
   Returns:
     The scan function with the signature ``(module, carry, *xs) -> (carry,
@@ -1298,6 +1301,7 @@ def scan(
     length=length,
     reverse=reverse,
     unroll=unroll,
+    _split_transpose=_split_transpose,
     data_transform=data_transform,
     metadata_params=metadata_params,
     methods=methods,