Merge pull request jax-ml#27809 from mattjj:26621

PiperOrigin-RevId: 745212009

Author: Google-ML-Automation

jax/experimental/shard_map.py

@@ -1145,6 +1145,14 @@ def _pbroadcast_batcher(vals_in, dims_in, *, axes, axis_index_groups):
 register_standard_check = \
     lambda prim: _check_rules.setdefault(prim, partial(_standard_check, prim))
 
+def _eq_rep(mesh, r1, r2) -> bool:
+  if r1 != r2 and r1 is None or r2 is None:
+    r1, r2 = _remove_none_rep(mesh, r1), _remove_none_rep(mesh, r2)
+  return r1 == r2
+
+def _remove_none_rep(mesh, r):
+  return set(mesh.axis_names) if r is None else r
+
 def _no_rewrite(prim, rule, mesh, in_rep, *args, **params):
   out_vals = prim.bind(*args,**params)
   out_rep = rule(mesh, *in_rep, **params)
@@ -1371,7 +1379,7 @@ def _scan_check(mesh, *in_rep, jaxpr, num_consts, num_carry, **_):
   _, carry_rep_in, _ = split_list(in_rep, [num_consts, num_carry])
   out_rep = _check_rep(mesh, jaxpr.jaxpr, in_rep)
   carry_rep_out, _ = split_list(out_rep, [num_carry])
-  if carry_rep_in != carry_rep_out:
+  if not all(map(partial(_eq_rep, mesh), carry_rep_in, carry_rep_out)):
     raise Exception("Scan carry input and output got mismatched replication "
                     f"types {carry_rep_in} and {carry_rep_out}. Please open an "
                     "issue at https://github.com/jax-ml/jax/issues, and as a "
@@ -1447,7 +1455,7 @@ def _cond_rule(mesh, *in_rep, branches):
   out_rep = _check_rep(mesh, branches[0].jaxpr, args_rep)
   for branch in branches[1:]:
     out_rep_ = _check_rep(mesh, branch.jaxpr, args_rep)
-    if out_rep_ != out_rep:
+    if not all(map(partial(_eq_rep, mesh), out_rep, out_rep_)):
       raise Exception("The branches of cond produced mismatched replication "
                       "types. Please open an issue at "
                       "https://github.com/jax-ml/jax/issues, and as a "
@@ -2189,7 +2197,7 @@ def _efficient_transpose_rewrite(fun, mesh, in_names, out_names_thunk):
 def _efficient_transpose_rewrite_nomatch(f, store, mesh, in_reps, *args):
   with core.take_current_trace() as parent:
     tag = core.TraceTag()
-    t = RewriteTrace(parent_trace = parent, tag = tag, mesh=mesh)
+    t = RewriteTrace(parent_trace=parent, tag=tag, mesh=mesh)
     in_tracers = map(partial(RewriteTracer, t), in_reps, args)
     with core.set_current_trace(t):
       ans = f(*in_tracers)

tests/shard_map_test.py

@@ -2798,6 +2798,30 @@ def f(x):
 
     f(x)  # doesn't crash
 
+  def test_rep_none_canonicalization(self):
+    # https://github.com/jax-ml/jax/issues/26621
+    if config.use_shardy_partitioner.value:
+      self.skipTest('complex values fail under shardy')
+    N = 8
+    xs = jnp.ones((8, N), dtype=jnp.int32)
+    variables = jax.random.normal(jax.random.key(1), (N, N), jnp.complex64)
+    mesh = jtu.create_mesh((2,), ('i',))
+    in_specs = (P(), P("i"),)
+    out_specs = P("i")
+
+    variables = jax.lax.with_sharding_constraint(variables, NamedSharding(mesh, P()))
+    xs = jax.lax.with_sharding_constraint(xs, NamedSharding(mesh, P('i')))
+
+    def fun(v, xs):
+        # Commenting this single line below makes everything work
+        v = jax.scipy.linalg.expm(v)
+        v = v.sum()
+        return v * xs.sum(axis=-1).astype(v.dtype)
+
+    res = fun(variables, xs)
+    fun_shard_map = shard_map(fun, mesh=mesh, in_specs=in_specs, out_specs=out_specs)
+    res = fun_shard_map(variables, xs)  # don't crash
+
 
 class FunSpec(NamedTuple):
   name: str