Fix shard_map check_rep for multiple outputs.

Author: justinjfu

jax/experimental/shard_map.py

@@ -658,7 +658,8 @@ def write(v: core.Var, val: RepType) -> None:
     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 isinstance(out_rep, (set, frozenset)):
+        out_rep = [out_rep] * len(e.outvars)
       foreach(write, e.outvars, out_rep)
     else:
       write(e.outvars[0], out_rep)
@@ -942,7 +943,7 @@ def to_val_rep_pair(self, val):
       raise Exception(f"Shouldn't have any non-shard_map tracers: {val}")
     else:
       val_ = _unmatch_spec(self.mesh, {}, val, self.context_mesh)
-      return val_, None
+      return val_, frozenset(self.mesh.axis_names) - self.auto
 
   def process_primitive(self, prim, tracers, params):
     in_vals, in_rep = unzip2(map(self.to_val_rep_pair, tracers))
@@ -958,7 +959,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 isinstance(out_rep, (set, frozenset)):
+        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,19 @@ 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",))
+    y = shard_map(
+          lambda x: jnp.split(x.squeeze(), 2),
+          mesh=mesh,
+          in_specs=(None,),
+          out_specs=P("i"),
+    )(xs)
+    expected = jnp.repeat(xs, 2, axis=0).reshape(2, 2, 10)
+    self.assertArraysEqual(y, expected)
+
   def test_vmap_spmd_axis_name(self):
     mesh = jtu.create_mesh((2, 2), ('x', 'y'))