Merge pull request jax-ml#27797 from mattjj:shmap-fix

PiperOrigin-RevId: 744827753

Author: Google-ML-Automation

jax/experimental/shard_map.py

@@ -958,7 +958,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
+      out_rep = (out_rep if isinstance(out_rep, (list, tuple))
+                 else [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'))