Repository: google/jax
Issue: Does shmap work with jax.lax.scan inside the mapped function?

### Description

I am attempting to perform a sharded_map of a function that has an inner jax.lax.scan (via `fori_loop`). When I run the shmapped function, the result is an error complaining that there is `No replication rule for scan`.

For further context, I am attempting to switch from `pmap` to `shmap`, as I would like to perform an outer `jit` on a function that calls the pmapped function internally. Inside this function, I call pmap over an axis named `'devices'` and the output of the pmapped function is further processed by other methods (some of which will be jitted as a result of the outer jit). However, the result is an error that states the `'devices'` axes is not available. My understanding is that this is due to pjit and pmap using fundamentally incompatible axes abstractions.

### What jax/jaxlib version are you using?

4.8

### Which accelerator(s) are you using?

GPU

### Additional system info

_No response_

### NVIDIA GPU info

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.141.03   Driver Version: 470.141.03   CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Quadro GP100        Off  | 00000000:81:00.0 Off |                    0 |
| 26%   38C    P0    30W / 235W |      0MiB / 16278MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  Quadro GP100        Off  | 00000000:82:00.0 Off |                    0 |
| 26%   36C    P0    30W / 235W |      0MiB / 16278MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

Changed files:
jax/experimental/shard_map.py, tests/shard_map_test.py

diff --git a/jax/experimental/shard_map.py b/jax/experimental/shard_map.py
index 49495ae19..6ab93d5b7 100644
--- a/jax/experimental/shard_map.py
+++ b/jax/experimental/shard_map.py
@@ -39,10 +39,10 @@ from jax._src import traceback_util
 from jax._src import util
 from jax._src.core import Tracer
 from jax._src.lax import (lax, parallel as lax_parallel, slicing,
-                          windowed_reductions, fft, linalg)
+                          windowed_reductions, fft, linalg, control_flow)
 from jax._src.util import (HashableFunction, HashablePartial, unzip2, unzip3,
                            as_hashable_function, memoize, partition_list,
-                           merge_lists)
+                           merge_lists, split_list)
 from jax.api_util import flatten_fun_nokwargs, shaped_abstractify
 from jax._src.interpreters import batching
 from jax._src.interpreters import mlir
@@ -133,7 +133,7 @@ SpecErrorType = enum.Enum('SpecErrorType', ['input', 'out'])
 def _check_specs(error_type: SpecErrorType, specs: Any) -> None:
   if error_type == SpecErrorType.input and specs is None:
     raise TypeError(
-        f"shard_map in_specs argument must be a pytree of "
+        "shard_map in_specs argument must be a pytree of "
         "`jax.sharding.PartitionSpec` instances, but it was None.\\n"
         "Instead of `in_specs=None`, did you mean `in_specs=P()`, "
         "where `P = jax.sharding.PartitionSpec`?")
@@ -712,6 +712,20 @@ def _pjit_rule(mesh, *in_rep, jaxpr, **kwargs):
 def _debug_callback_rule(mesh, *in_rep, **_):
   return []
 
+@register_rule(control_flow.loops.scan_p)
+def _scan_rule(mesh, *in_rep, jaxpr, num_consts, num_carry, linear, length,
+               reverse, unroll):
+  const_rep, carry_rep, xs_rep = split_list(in_rep, [num_consts, num_carry])
+  for _ in range(1 + num_carry):
+    out_rep = _output_rep(mesh, jaxpr.jaxpr, [*const_rep, *carry_rep, *xs_rep])
+    if carry_rep == out_rep[:num_carry]:
+      break
+    else:
+      carry_rep = map(op.and_, carry_rep, out_rep[:num_carry])
+  else:
+    assert False, 'Fixpoint not reached'
+  return out_rep
+
 # Batching
 
 def _shard_map_batch(
@@ -831,7 +845,7 @@ def _shard_map_partial_eval(trace, shard_map_p, f, tracers, mesh, in_names,
                       out_names_thunk=known_out_names, check_rep=check_rep)
   out = shard_map_p.bind(f_known, *in_consts, **known_params)
   out_knowns, out_avals_sharded, jaxpr, env = aux()
-  out_consts, res = pe.split_list(out, [len(out) - len(jaxpr.constvars)])
+  out_consts, res = split_list(out, [len(out) - len(jaxpr.constvars)])
   with core.extend_axis_env_nd(mesh.shape.items()):
     jaxpr = pe.convert_constvars_jaxpr(jaxpr)
   unk_out_names, _ = pe.partition_list(out_knowns, out_names_thunk())
@@ -865,7 +879,7 @@ def _shard_map_partial_eval_post_process(
 
   def todo(out):
     trace = main.with_cur_sublevel()
-    out_consts, res = pe.split_list(out, [len(out) - len(jaxpr.constvars)])
+    out_consts, res = split_list(out, [len(out) - len(jaxpr.constvars)])
     const_tracers = map(trace.new_instantiated_const, res)
     env_tracers = map(trace.full_raise, env)
 
diff --git a/tests/shard_map_test.py b/tests/shard_map_test.py
index 2d7a39646..7ee0ac56b 100644
--- a/tests/shard_map_test.py
+++ b/tests/shard_map_test.py
@@ -595,6 +595,50 @@ class ShardMapTest(jtu.JaxTestCase):
 
     shard_map(f, mesh, in_specs=P(), out_specs=P())(3.)  # doesn't crash
 
+  def test_scan_rep_rule(self):
+    mesh = jtu.create_global_mesh((2, 2,), ('x', 'y'))
+
+    def f(x, y, z):
+      x, y, z = x.sum(), y.sum(), z.sum()
+      def body(c, _):
+        c, *cs = c
+        return (*cs, c), None
+      out, _  = jax.lax.scan(body, (x, y, z), None, length=3)
+      return [jnp.expand_dims(a, 0) for a in out]
+
+    x = jnp.arange(4)
+
+    # doesn't crash, because out_spec assumes no replication (and there is none)
+    shard_map(f, mesh, in_specs=(P(None), P('x'), P(('x', 'y'))),
+              out_specs=P(('x', 'y')))(x, x, x)
+
+    # does crash, because output incorrectly promises replication
+    with self.assertRaisesRegex(ValueError, "require replication"):
+      shard_map(f, mesh, in_specs=(P(None), P('x'), P(('x', 'y'))),
+                out_specs=P('x'))(x, x, x)
+    with self.assertRaisesRegex(ValueError, "require replication"):
+      shard_map(f, mesh, in_specs=(P(None), P('x'), P(('x', 'y'))),
+                out_specs=P('y'))(x, x, x)
+    with self.assertRaisesRegex(ValueError, "require replication"):
+      shard_map(f, mesh, in_specs=(P(None), P('x'), P(('x', 'y'))),
+                out_specs=P(None))(x, x, x)
+
+    def g(x, y, z):
+      x, y, z = x.sum(), y.sum(), z.sum()
+      def body(c, _):
+        return c, None
+      out, _  = jax.lax.scan(body, (x, y, z), None, length=1)
+      return [jnp.expand_dims(a, 0) for a in out]
+
+    # doesn't crash, because everything matches
+    shard_map(g, mesh, in_specs=(P(None), P('x'), P(('x', 'y'))),
+              out_specs=[P(None), P('x'), P(('x', 'y'))])(x, x, x)
+
+    # does crash, because the second guy is wrong
+    with self.assertRaisesRegex(ValueError, "require replication"):
+      shard_map(g, mesh, in_specs=(P(None), P('x'), P(('x', 'y'))),
+                out_specs=[P(None), P(None), P(('x', 'y'))])(x, x, x)
+
 
 class FunSpec(NamedTuple):
   name: str