Repository: google/jax

Issue:
BUG: `shard_alike` requires newer jaxlib
This failure occurs when running jax at HEAD with jaxlib 0.4.23
```python
import os
os.environ["XLA_FLAGS"] = " --xla_force_host_platform_device_count=8"

import jax
import jaxlib
import jax.numpy as jnp
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

print(f"{jax.__version__=}")
print(f"{jaxlib.__version__=}")

x = jnp.arange(16, dtype='float32').reshape(8, 2)
mesh = Mesh(jax.devices(), axis_names=("i"))

x = jax.device_put(x, NamedSharding(mesh, P('i', None)))
print(x.sharding)

y = jax.numpy.zeros_like(x)
print(y.sharding)
```
```pytb
jax.__version__='0.4.24.dev20240117'
jaxlib.__version__='0.4.23'
NamedSharding(mesh=Mesh('i': 8), spec=PartitionSpec('i', None))
---------------------------------------------------------------------------
XlaRuntimeError                           Traceback (most recent call last)
[<ipython-input-1-c390646586f4>](https://localhost:8080/#) in <cell line: 18>()
     16 print(x.sharding)
     17 
---> 18 y = jax.numpy.zeros_like(x)
     19 print(y.sharding)

8 frames
    [... skipping hidden 10 frame]

[/usr/local/lib/python3.10/dist-packages/jax/_src/interpreters/pxla.py](https://localhost:8080/#) in __call__(self, *args)
   1154       self._handle_token_bufs(result_token_bufs, sharded_runtime_token)
   1155     else:
-> 1156       results = self.xla_executable.execute_sharded(input_bufs)
   1157     if dispatch.needs_check_special():
   1158       out_arrays = results.disassemble_into_single_device_arrays()

XlaRuntimeError: INVALID_ARGUMENT: Executable expected parameter 1 of size 8 but got buffer with incompatible size 64: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).
```
It can be reproduced more directly by replacing the last line with this:
```python
jax.experimental.shard_alike.shard_alike(jnp.zeros(x.shape), x)
```
It looks like this bug was introduced to `full_like` in #19115.


Discussion:

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -62,6 +62,7 @@
   standard_primitive)
 from jax._src import xla_bridge
 from jax._src.lib import xla_client
+from jax._src.lib import xla_extension_version
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import chlo
 from jax._src.lib.mlir.dialects import hlo
@@ -1375,7 +1376,14 @@ def full_like(x: ArrayLike | DuckTypedArray,
   # TODO(yashkatariya): Use shard_like in tracing mode too i.e. remove the
   # ArrayImpl check.
   if shape is None and isinstance(x, array.ArrayImpl):
-    return shard_alike.shard_alike(x, val)[1]
+    if xla_extension_version < 227:
+      sharding = x.sharding  # type: ignore[union-attr]
+      if (not dispatch.is_single_device_sharding(sharding) and
+          not isinstance(sharding, PmapSharding)):
+        return array.make_array_from_callback(
+            type_cast(array.Shape, fill_shape), sharding, lambda idx: val[idx])
+    else:
+      return shard_alike.shard_alike(x, val)[1]
   return val
 
 
diff --git a/jax/_src/shard_alike.py b/jax/_src/shard_alike.py
--- a/jax/_src/shard_alike.py
+++ b/jax/_src/shard_alike.py
@@ -23,6 +23,7 @@
 from jax._src.interpreters import batching
 from jax._src.util import safe_zip
 from jax._src.lib import xla_client as xc
+from jax._src.lib import xla_extension_version
 from jax._src.api_util import shaped_abstractify
 from jax._src.lib.mlir import ir
 
@@ -30,6 +31,8 @@
 
 def shard_alike(x, y):
   """Shards x and y alike."""
+  if xla_extension_version < 227:
+    raise ValueError("shard_alike requires jaxlib v0.4.24 or newer.")
   x_flat, x_tree = tree_flatten(x)
   y_flat, y_tree = tree_flatten(y)
 


Test patch:
diff --git a/tests/shard_alike_test.py b/tests/shard_alike_test.py
--- a/tests/shard_alike_test.py
+++ b/tests/shard_alike_test.py
@@ -50,6 +50,16 @@ def tearDownModule():
   xla_bridge.get_backend.cache_clear()
 
 
+class ShardAlikeDownstreamTest(jtu.JaxTestCase):
+
+  def test_full_like(self):
+    x = jnp.arange(16, dtype='float32').reshape(8, 2)
+    mesh = jtu.create_global_mesh((8,), ("i",))
+    x = jax.device_put(x, NamedSharding(mesh, P('i', None)))
+    y = jnp.full_like(x, 1)
+    self.assertEqual(x.sharding, y.sharding)
+
+
 class ShardAlikeTest(jtu.JaxTestCase):
 
   def setUp(self):