Fix shard_map check_rep for multiple outputs.

Author: justinjfu

jax/experimental/shard_map.py

@@ -657,8 +657,10 @@ def write(v: core.Var, val: RepType) -> None:
   for e in jaxpr.eqns:
     rule = _check_rules.get(e.primitive, partial(_rule_missing, e.primitive))
     out_rep = rule(mesh, *map(read, e.invars), **e.params)
+
     if e.primitive.multiple_results:
-      out_rep = [out_rep] * len(e.outvars) if type(out_rep) is set else out_rep
+      if type(out_rep) is set or out_rep is None:
+        out_rep = [out_rep] * len(e.outvars)
       foreach(write, e.outvars, out_rep)
     else:
       write(e.outvars[0], out_rep)
@@ -958,7 +960,8 @@ def process_primitive(self, prim, tracers, params):
     rep_rule = _check_rules.get(prim, partial(_rule_missing, prim))
     out_rep = rep_rule(self.mesh, *in_rep, **params) if self.check else set()
     if prim.multiple_results:
-      out_rep = [out_rep] * len(out_vals) if type(out_rep) is set else out_rep
+      if type(out_rep) is set or out_rep is None:
+        out_rep = [out_rep] * len(out_vals)
       return map(partial(ShardMapTracer, self), out_rep, out_vals)
     return ShardMapTracer(self, out_rep, out_vals)
 

tests/shard_map_test.py

@@ -685,6 +685,17 @@ def f3():
     f3()
     jax.jit(f3)()
 
+  def test_multiple_result_primitive_with_none_sharding(self):
+    # https://github.com/jax-ml/jax/issues/27673
+    xs = jnp.arange(20).reshape(2, 10)
+    mesh = jtu.create_mesh((2,), ("i",))
+    shard_map(
+          lambda x: jnp.split(x.squeeze(), 2),
+          mesh=mesh,
+          in_specs=(None,),
+          out_specs=P("i"),
+    )(xs)
+
   def test_vmap_spmd_axis_name(self):
     mesh = jtu.create_mesh((2, 2), ('x', 'y'))