Convert shardings in jit path to OpShardingSharding to avoid recomp…

…ilation when semantically similar shardings are used in `jit`. PiperOrigin-RevId: 477626548
google · Sep 29, 2022 · 163b7e2 · 163b7e2
1 parent 500f8b7
commit 163b7e2
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/jax/_src/api.py b/jax/_src/api.py
@@ -670,13 +670,15 @@ def _jit_lower(fun, static_argnums, static_argnames, device, backend,
 
   def arg_spec(x):
     from jax._src.sharding import PmapSharding
+    from jax.experimental import pjit
     # like xla.arg_spec but duck-types on x.shape and x.dtype
     aval = None if jax.config.jax_dynamic_shapes else shaped_abstractify(x)
     if jax.config.jax_array:
       if hasattr(x, 'sharding'):
         if isinstance(x.sharding, PmapSharding):
           return aval, None
-        return aval, (x.sharding if x._committed else None)
+        return aval, (pjit.to_op_sharding_sharding(x.sharding, x.ndim)
+                      if x._committed else None)
       else:
         return aval, None
     else:

diff --git a/jax/_src/dispatch.py b/jax/_src/dispatch.py
@@ -91,13 +91,15 @@
 
 def arg_spec(x: Any) -> ArgSpec:
   from jax._src.sharding import PmapSharding
+  from jax.experimental import pjit
 
   aval = xla.abstractify(x)
   try:
     if config.jax_array:
       if isinstance(x.sharding, PmapSharding):
         return aval, None
-      return aval, (x.sharding if x._committed else None)
+      return aval, (pjit.to_op_sharding_sharding(x.sharding, x.ndim)  # type: ignore
+                    if x._committed else None)
     else:
       return aval, x._device
   except: