nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: jax-ml/jax
Pull request #25930 (2025-01-16 19:00:32)

Allow resharding between tokens on a single device
Allow resharding between tokens on a single device
and multiple devices.

Whenever this happens we can essentially introduce an effects barrier
instead of doing the normal device -> host -> device transfer.

Fixes https://github.com/jax-ml/jax/issues/25671.



Patch:
diff --git a/jax/_src/array.py b/jax/_src/array.py
index 2ee8b01c7..802a3f5f6 100644
--- a/jax/_src/array.py
+++ b/jax/_src/array.py
@@ -1205,8 +1205,12 @@ pxla.local_result_handlers[core.ShapedArray] = _array_local_result_handler
 # Token handlers
 
 def _token_shard_arg(xs, shardings, layouts, copy_semantics):
-  return _array_shard_arg([x._buf for x in xs], shardings, layouts,
-                          copy_semantics)
+  results = []
+  for x, sharding, layout in safe_zip(xs, shardings, layouts):
+    x.block_until_ready()
+    x = np.array([], dtype=bool)
+    results.append(api.device_put(x, Layout(layout, sharding)))
+  return results
 pxla.shard_arg_handlers[core.Token] = _token_shard_arg
 
 


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 92bbf57ef..6abc83225 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -4327,6 +4327,30 @@ class ArrayPjitTest(jtu.JaxTestCase):
         out_specs=P(None, 'i')))
     f(jnp.zeros((2, 16)), jnp.ones(2))
 
+  def test_empty_io_callback_under_shard_map_reshard_to_singledev(self):
+    if config.use_shardy_partitioner.value:
+      self.skipTest("Shardy errors out on empty callbacks.")
+    mesh = jtu.create_mesh((4,), 'i')
+
+    def empty_callback(x):
+      return
+
+    def _f(x, y):
+      jax.experimental.io_callback(
+          empty_callback, (), x, ordered=True)
+      return x + y[..., jnp.newaxis]
+
+    f = jax.jit(shard_map(
+        _f, mesh, in_specs=(P(None, 'i'), P(None)),
+        out_specs=P(None, 'i')))
+    f(jnp.zeros((2, 16)), jnp.ones(2))
+
+    _f(jnp.zeros((2, 16)), jnp.ones(2))
+
+    jax.jit(_f)(jnp.zeros((2, 16)), jnp.ones(2))
+
+    jax.effects_barrier()
+
   def test_jit_trace_lower_and_compile(self):
     def f(x):
       return x * 2
Repository: jax-ml/jax
Pull request #26995 (2025-03-07 17:25:28)

[sharding_in_types] Fix a bug where `empty_array` in scan was created with the wrong spec when `unroll > 1`.
[sharding_in_types] Fix a bug where `empty_array` in scan was created with the wrong spec when `unroll > 1`.



Patch:
diff --git a/jax/_src/lax/control_flow/loops.py b/jax/_src/lax/control_flow/loops.py
index caabace86..5179da751 100644
--- a/jax/_src/lax/control_flow/loops.py
+++ b/jax/_src/lax/control_flow/loops.py
@@ -450,7 +450,7 @@ def _scan_impl(*args, reverse, length, num_consts, num_carry, jaxpr, linear,
     num_trips, remainder = 0, length
   if unroll == 1:
     xss = xs_
-    yss = _map(partial(_empty_array, (length,), None), y_avals)
+    yss = _map(partial(_empty_array, (length,), (None,)), y_avals)
   else:
     if remainder:
       if not reverse:
@@ -458,7 +458,7 @@ def _scan_impl(*args, reverse, length, num_consts, num_carry, jaxpr, linear,
       else:
         xs_rem, xs_ = unzip2(_map(partial(_split_leading, remainder), xs_))
     xss = [lax.reshape(x, (num_trips, unroll, *x.shape[1:])) for x in xs_]
-    yss = _map(partial(_empty_array, (num_trips, unroll), None), y_avals)
+    yss = _map(partial(_empty_array, (num_trips, unroll), (None, None)), y_avals)
 
   def inner(n, carry, xs):
     ys = []
@@ -509,7 +509,7 @@ def _split_leading(sz, x):
 def _concat(a, b): return lax.concatenate([a, b], 0)
 
 def _empty_array(prefix, length_spec, aval):
-  sharding = aval.sharding.with_spec((length_spec, *aval.sharding.spec))
+  sharding = aval.sharding.with_spec((*length_spec, *aval.sharding.spec))
   return lax.broadcast(lax.empty(aval.dtype), (*prefix, *aval.shape),
                        out_sharding=sharding)
 


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index b0a59027f..720a77410 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -7080,6 +7080,20 @@ class ShardingInTypesTest(jtu.JaxTestCase):
         ValueError, 'does not evenly divide the dimension size'):
       core.ShapedArray((5, 2), jnp.int32, sharding=s)
 
+  @jtu.with_user_mesh((2, 2), ('x', 'y'))
+  def test_scan_unroll(self, mesh):
+    np_inp = np.arange(64, dtype=jnp.float32).reshape(8, 8)
+    arr = jax.device_put(np_inp, NamedSharding(mesh, P(None, 'y')))
+    carry = jnp.ones((8,), dtype=jnp.float32)
+
+    @jax.jit
+    def f(carry, xs):
+      def body(carry, x):
+        return carry + x, x
+      return jax.lax.scan(body, carry, xs, unroll=2)
+
+    f(carry, arr)  # doesn't crash
+
 
 @jtu.pytest_mark_if_available('multiaccelerator')
 class PJitErrorTest(jtu.JaxTestCase):
Repository: jax-ml/jax
Pull request #27000 (2025-03-07 21:29:47)

[Pallas] Small fix to TPU interpret mode (input_output_aliases + scalar args).
[Pallas] Small fix to TPU interpret mode (input_output_aliases + scalar args).



Patch:
diff --git a/jax/_src/pallas/mosaic/interpret.py b/jax/_src/pallas/mosaic/interpret.py
index 4f2d0b4f3..e2086d6af 100644
--- a/jax/_src/pallas/mosaic/interpret.py
+++ b/jax/_src/pallas/mosaic/interpret.py
@@ -1346,7 +1346,9 @@ def interpret_pallas_call(
   output_buffer_ids = []
   output_buffer_shapes = []
   output_vals = _initialize_output_vals(
-      grid_mapping.block_mappings_output, args, input_output_aliases)
+      grid_mapping.block_mappings_output,
+      scalars + input_args,
+      input_output_aliases)
   num_outputs = grid_mapping.num_outputs
   output_block_shapes = block_shapes[num_inputs : num_inputs + num_outputs]
   for out_val, bs in zip(output_vals, output_block_shapes):


Test patch:
diff --git a/tests/pallas/tpu_pallas_interpret_test.py b/tests/pallas/tpu_pallas_interpret_test.py
index 632dba949..1219f37fb 100644
--- a/tests/pallas/tpu_pallas_interpret_test.py
+++ b/tests/pallas/tpu_pallas_interpret_test.py
@@ -68,25 +68,30 @@ class InterpretTest(jtu.JaxTestCase):
     z = matmul(x, y)
     np.testing.assert_allclose(z, x @ y, atol=1e-4)
 
-  def test_dynamic_grid(self):
-    def kernel(x_ref, o_ref):
-      o_ref[...] = x_ref[...]
+  def test_dynamic_grid_and_aliasing(self):
+    def kernel(s_ref, x_ref, o_ref):
+      o_ref[...] = x_ref[...] + s_ref[0].astype(x_ref.dtype)
 
     iters = jax.random.randint(jax.random.key(0), (), 10, 20, dtype=jnp.int32)
     @jax.jit
-    def f(x):
+    def f(s, x):
       return pl.pallas_call(
           kernel,
           out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
           grid=(iters,),
-          in_specs=(pl.BlockSpec(x.shape, lambda i: (0, 0)),),
+          in_specs=[
+              pl.BlockSpec(memory_space=pltpu.TPUMemorySpace.SMEM),
+              pl.BlockSpec(x.shape, lambda i: (0, 0)),
+          ],
           out_specs=pl.BlockSpec(x.shape, lambda i: (0, 0)),
+          input_output_aliases={1: 0},
           interpret=mosaic_interpret.TPUInterpretParams()
-      )(x)
+      )(s, x)
 
+    s = jnp.array([1], dtype=jnp.int32)
     x = jnp.arange(32 * 128.).reshape((32, 128))
-    y = f(x)
-    np.testing.assert_allclose(y, x)
+    y = f(s, x)
+    np.testing.assert_allclose(y, x + 1.0)
 
   @parameterized.parameters('eager', 'on_wait')
   def test_race_detection(self, dma_execution_mode):
Repository: jax-ml/jax
Pull request #27011 (2025-03-08 01:59:31)

Better error message when `raise_if_error()` is called within a traced context
Better error message when `raise_if_error()` is called within a traced context



Patch:
diff --git a/jax/_src/error_check.py b/jax/_src/error_check.py
index edfcdf3f7..11e65a7dd 100644
--- a/jax/_src/error_check.py
+++ b/jax/_src/error_check.py
@@ -88,12 +88,17 @@ def raise_if_error() -> None:
   """Raise error if an error is set.
 
   This function should be called after the computation is finished. It should
-  be used outside jit.
+  not be called within a traced context, such as within a jitted function."
   """
   if _error_storage.ref is None:  # if not initialized, do nothing
     return
 
   error_code = _error_storage.ref[...]
+  if isinstance(error_code, core.Tracer):
+    raise ValueError(
+        "raise_if_error() should not be called within a traced context, such as"
+        " within a jitted function."
+    )
   if error_code == jnp.uint32(_NO_ERROR):
     return
   _error_storage.ref[...] = jnp.uint32(_NO_ERROR)


Test patch:
diff --git a/tests/error_check_test.py b/tests/error_check_test.py
index 653f901a6..5cdde30b1 100644
--- a/tests/error_check_test.py
+++ b/tests/error_check_test.py
@@ -170,6 +170,26 @@ class ErrorCheckTests(jtu.JaxTestCase):
     _ = body(init, xs)
     error_check.raise_if_error()  # should not raise error
 
+  @parameterized.product(jit=[True, False])
+  def test_raise_if_error_fails_in_traced_context(self, jit):
+    def f(x):
+      error_check.set_error_if(x <= 0, "x must be greater than 0")
+      return x + 1
+
+    if jit:
+      f = jax.jit(f)
+
+    x = jnp.full((4,), 1, dtype=jnp.int32)
+    f(x)
+    with self.assertRaises(
+        ValueError,
+        msg=(
+            "raise_if_error() should not be called within a traced context,"
+            " such as within a jitted function."
+        ),
+    ):
+      jax.jit(error_check.raise_if_error)()
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())
Repository: jax-ml/jax
Pull request #27078 (2025-03-11 22:53:08)

Rename `get_ty` to `typeof` which is an alias of `get_aval`
Rename `get_ty` to `typeof` which is an alias of `get_aval`



Patch:
diff --git a/jax/__init__.py b/jax/__init__.py
index 950c3ed4b..ae3bac4ad 100644
--- a/jax/__init__.py
+++ b/jax/__init__.py
@@ -79,7 +79,7 @@ from jax._src.lib import xla_client as _xc
 Device = _xc.Device
 del _xc
 
-from jax._src.core import get_ty as get_ty
+from jax._src.core import typeof as typeof
 from jax._src.api import effects_barrier as effects_barrier
 from jax._src.api import block_until_ready as block_until_ready
 from jax._src.ad_checkpoint import checkpoint_wrapper as checkpoint  # noqa: F401
diff --git a/jax/_src/core.py b/jax/_src/core.py
index 9d8edeb8b..b17e26255 100644
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -1576,7 +1576,7 @@ def get_aval(x):
     return get_aval(x.__jax_array__())
   raise TypeError(f"Argument '{x}' of type '{typ}' is not a valid JAX type")
 
-get_ty = get_aval
+typeof = get_aval
 
 def is_concrete(x):
   return to_concrete_value(x) is not None


Test patch:
diff --git a/tests/mutable_array_test.py b/tests/mutable_array_test.py
index c510c2cfa..4c6a8eb7a 100644
--- a/tests/mutable_array_test.py
+++ b/tests/mutable_array_test.py
@@ -216,8 +216,8 @@ class MutableArrayTest(jtu.JaxTestCase):
 
     @jax.jit
     def f(x_ref):
-      self.assertEqual(core.get_ty(x_ref).sharding.spec,
-                       core.get_ty(x_ref[...]).sharding.spec)
+      self.assertEqual(core.typeof(x_ref).sharding.spec,
+                       core.typeof(x_ref[...]).sharding.spec)
       y = x_ref[...] + 1
       return y
 
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index bd7954d60..4cd1af9d3 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -4883,11 +4883,11 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     arr = jax.device_put(np_inp, s)
 
     def f(x):
-      self.assertEqual(jax.get_ty(x).sharding.spec, s.spec)
+      self.assertEqual(jax.typeof(x).sharding.spec, s.spec)
       x = x * 2
-      self.assertEqual(jax.get_ty(x).sharding.spec, s.spec)
+      self.assertEqual(jax.typeof(x).sharding.spec, s.spec)
       x = x * x
-      self.assertEqual(jax.get_ty(x).sharding.spec, s.spec)
+      self.assertEqual(jax.typeof(x).sharding.spec, s.spec)
       return x
 
     # Eager mode
Repository: jax-ml/jax
Pull request #27100 (2025-03-12 19:33:00)

[sharding_in_types] Make `reshard` work with np.array.
[sharding_in_types] Make `reshard` work with np.array.



Patch:
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index 06892aa9f..b2fb692e6 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -2717,11 +2717,6 @@ def _mesh_cast_abstract_eval(aval, dst_sharding):
         f' axis_types={src_sharding.mesh.axis_types} and dst'
         f' axis_types={dst_sharding.mesh.axis_types}. To reshard between the'
         ' same mesh, use `jax.sharding.reshard` instead?')
-  if len(src_sharding.spec) != len(dst_sharding.spec):
-    raise ValueError(
-        'Length of source sharding spec should be equal to destination'
-        f' sharding spec. Got source spec={src_sharding.spec} and destination'
-        f' spec={dst_sharding.spec}')
   if src_sharding.mesh._any_axis_explicit and dst_sharding.mesh._any_axis_explicit:
     for s, d in safe_zip(flatten_spec(src_sharding.spec),
                          flatten_spec(dst_sharding.spec)):
@@ -2786,7 +2781,8 @@ reshard_p = core.Primitive('reshard')
 
 def _reshard_abstract_eval(aval, dst_sharding):
   src_sharding = aval.sharding
-  if src_sharding.mesh.abstract_mesh != dst_sharding.mesh.abstract_mesh:
+  if (not src_sharding.mesh.empty and
+      src_sharding.mesh.abstract_mesh != dst_sharding.mesh.abstract_mesh):
     raise ValueError(
         f'Mesh of the input {src_sharding.mesh.abstract_mesh} does not'
         ' equal the mesh of the target sharding'


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 4cd1af9d3..53990e79e 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -6624,6 +6624,10 @@ class ShardingInTypesTest(jtu.JaxTestCase):
     out = hf(arr)
     self.assertEqual(out.sharding, NamedSharding(mesh, P('x', 'y')))
 
+    hf = auto_axes(f, axes=('x', 'y'), out_shardings=P('x'))
+    out = hf(arr)
+    self.assertEqual(out.sharding, NamedSharding(mesh, P('x', None)))
+
   @jtu.with_user_mesh((2, 2), ('x', 'y'),
                       axis_types={AxisTypes.Auto: ('x', 'y')})
   def test_full_visible_outside_jit(self, mesh):
@@ -7068,6 +7072,17 @@ class ShardingInTypesTest(jtu.JaxTestCase):
 
     f(carry, arr)  # doesn't crash
 
+  @jtu.with_user_mesh((2,), ('x',))
+  def test_reshard_with_np_array(self, mesh):
+    out = reshard(np.arange(8), P('x'))
+    self.assertEqual(out.sharding, NamedSharding(mesh, P('x')))
+
+    @jax.jit
+    def f(x):
+      return reshard(x, P('x'))
+    out = f(np.arange(8))
+    self.assertEqual(out.sharding, NamedSharding(mesh, P('x')))
+
 
 @jtu.pytest_mark_if_available('multiaccelerator')
 class PJitErrorTest(jtu.JaxTestCase):
Repository: jax-ml/jax
Pull request #27163 (2025-03-14 16:04:17)

Change the way that batching.spec_types is updated.
Change the way that batching.spec_types is updated.

There's no reason why not two custom vmappable types cannot share the same spec_type. However, spec_types was a set, which can cause bugs / exceptions.

Suppose that I register two vmappable data_types sharing the same spec_type, and then unregister one of the two. Then, the spec_type is no longer in the set to support the second data_type. Also, an exception will be raised if I try to unregister the two vmappable types (the second call to spec_types.remove).

When unregistering a data type, instead of removing its spec_type from the set, we regenerate the set from the remaining vmappable types.



Patch:
diff --git a/jax/_src/interpreters/batching.py b/jax/_src/interpreters/batching.py
index 40dbe0018..03c9a9510 100644
--- a/jax/_src/interpreters/batching.py
+++ b/jax/_src/interpreters/batching.py
@@ -322,12 +322,15 @@ vmappables: dict[type, tuple[type, type]] = {}
 spec_types: set[type] = {JumbleAxis}
 
 def unregister_vmappable(data_type: type) -> None:
-  spec_type, axis_size_type = vmappables.pop(data_type)
-  spec_types.remove(spec_type)
+  _, axis_size_type = vmappables.pop(data_type)
   del to_elt_handlers[data_type]
   del from_elt_handlers[data_type]
   if axis_size_type in make_iota_handlers:
     del make_iota_handlers[axis_size_type]
+  global spec_types
+  spec_types = (
+      {JumbleAxis} | {spec_type for spec_type, _ in vmappables.values()}
+  )
 
 def is_vmappable(x: Any) -> bool:
   return type(x) is Jumble or type(x) in vmappables


Test patch:
diff --git a/tests/batching_test.py b/tests/batching_test.py
index bab18ce53..f2a4e8c34 100644
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -1356,6 +1356,32 @@ class VmappableTest(jtu.JaxTestCase):
       self.assertEqual(ans.names, expected.names)
       self.assertAllClose(ans.data, expected.data)
 
+  def test_types_with_same_spec(self):
+    # We register NamedArray.
+    batching.register_vmappable(NamedArray, NamedMapSpec, int,
+                              named_to_elt, named_from_elt, None)
+
+    # We then register another type that uses NamedMapSpec as the spec_type too,
+    # and immediately unregister it.
+    class Foo:
+      pass
+    batching.register_vmappable(Foo, NamedMapSpec, int,
+                                named_to_elt, named_from_elt, None)
+    batching.unregister_vmappable(Foo)
+
+    # We should still be able to use vmap on NamedArray.
+    def f(x):
+      return named_mul(x, x)
+
+    x = NamedArray(['i', 'j'], jnp.arange(12.).reshape(3, 4))
+    ans = jax.jit(f)(x)
+    expected = NamedArray(['i', 'j'], jnp.arange(12.).reshape(3, 4) ** 2)
+
+    self.assertEqual(ans.names, expected.names)
+    self.assertAllClose(ans.data, expected.data)
+
+    # And unregister NamedArray without exceptions.
+    batching.unregister_vmappable(NamedArray)
 
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())
Repository: jax-ml/jax
Pull request #27199 (2025-03-17 20:36:09)

Support tuples in custom_partitioning.
Support tuples in custom_partitioning.



Patch:
diff --git a/jax/_src/custom_partitioning.py b/jax/_src/custom_partitioning.py
index 658a6f7a2..537407151 100644
--- a/jax/_src/custom_partitioning.py
+++ b/jax/_src/custom_partitioning.py
@@ -179,7 +179,7 @@ def _custom_partitioning_partition(arg_shapes, arg_shardings, result_shape,
       for sharding, s in zip(result_shardings, result_shapes)
   ]
   closed_jaxpr = jax.make_jaxpr(lower_fn, axis_env=list(mesh.shape.items()))(
-      *tiled_args
+      *info.in_tree.unflatten(tiled_args)
   )
   if ([(o.shape, o.dtype) for o in closed_jaxpr.out_avals] !=
       [(t.shape, t.dtype) for t in tiled_results]):


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 76336920b..293b37a9f 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -1680,6 +1680,47 @@ class CustomPartitionerTest(jtu.JaxTestCase):
     jit_f = jax.jit(f, in_shardings=s, out_shardings=s)
     self.assertArraysEqual(x, jit_f(x))
 
+  @jtu.with_mesh([('x', 4), ('y', 2)])
+  def test_custom_partitioner_pytree_inputs(self):
+    self.skip_if_custom_partitioning_not_supported()
+
+    def partition(mesh, arg_shapes, result_shape):
+      def lower_fn(xs):
+        x, y, z = xs
+        return x + y + z
+
+      return (
+          mesh,
+          lower_fn,
+          arg_shapes[0][0].sharding,
+          jax.tree.map(lambda x: x.sharding, arg_shapes),
+      )
+
+    def infer_sharding_from_operands(mesh, arg_shapes, result_shape):
+      return arg_shapes[0][0].sharding
+
+    def propagate_user_sharding(mesh, user_shape):
+      return user_shape.sharding
+
+    @custom_partitioning
+    def f(xs):
+      x, y, z = xs
+      return x + y + z
+
+    f.def_partition(
+        infer_sharding_from_operands=infer_sharding_from_operands,
+        partition=partition,
+        propagate_user_sharding=propagate_user_sharding,
+        sharding_rule='i j, i j, i j -> i j',
+    )
+
+    def f2(a):
+      return a + f((a, a, a))
+
+    pjit_f = pjit(f2, in_shardings=(P(None, 'x')), out_shardings=P('x'))
+    x = np.asarray(np.random.randint(0, 20, (32, 16)), dtype=np.float32)
+    self.assertArraysEqual(x * 4, pjit_f(x))
+
 
 @jtu.pytest_mark_if_available('multiaccelerator')
 class AutoShardingPjitTest(jtu.JaxTestCase):
Repository: jax-ml/jax
Pull request #27305 (2025-03-20 19:48:09)

Allow `tuple(arrays)` as an input to `make_array_from_single_device_arrays`. Fixes https://github.com/jax-ml/jax/issues/27303
Allow `tuple(arrays)` as an input to `make_array_from_single_device_arrays`. Fixes https://github.com/jax-ml/jax/issues/27303



Patch:
diff --git a/jax/_src/array.py b/jax/_src/array.py
index e49963ccd..ee1960268 100644
--- a/jax/_src/array.py
+++ b/jax/_src/array.py
@@ -1024,7 +1024,7 @@ def make_array_from_single_device_arrays(
     shape : Shape of the output ``jax.Array``. This conveys information already included with
       ``sharding`` and ``arrays`` and serves as a double check.
     sharding: Sharding: A global Sharding instance which describes how the output jax.Array is laid out across devices.
-    arrays: Sequence of ``jax.Array``\s that are each single device addressable. ``len(arrays)``
+    arrays: `list` or `tuple` of ``jax.Array``\s that are each single device addressable. ``len(arrays)``
       must equal ``len(sharding.addressable_devices)`` and the shape of each array must be the same. For multiprocess code,
       each process will call with a different ``arrays`` argument that corresponds to that processes' data.
       These arrays are commonly created via ``jax.device_put``.
@@ -1071,14 +1071,15 @@ def make_array_from_single_device_arrays(
   if dtypes.issubdtype(aval.dtype, dtypes.extended):
     return aval.dtype._rules.make_sharded_array(aval, sharding, arrays,
                                                 committed=True)
+  arrays = list(arrays) if isinstance(arrays, tuple) else arrays
   # TODO(phawkins): ideally the cast() could be checked.
   try:
     return ArrayImpl(aval, sharding, cast(Sequence[ArrayImpl], arrays),
                     committed=True)
   except TypeError:
-    if not isinstance(arrays, Sequence):
+    if not isinstance(arrays, list):
       raise TypeError("jax.make_array_from_single_device_arrays `arrays` "
-                      "argument must be a Sequence (list or tuple), but got "
+                      "argument must be a list or tuple, but got "
                       f"{type(arrays)}.")
     if any(isinstance(arr, core.Tracer) for arr in arrays):
       raise ValueError(


Test patch:
diff --git a/tests/array_test.py b/tests/array_test.py
index cc8990828..6100283cc 100644
--- a/tests/array_test.py
+++ b/tests/array_test.py
@@ -1301,6 +1301,18 @@ class ShardingTest(jtu.JaxTestCase):
     with self.assertRaisesRegex(TypeError, msg):
       jax.jit(f)(x)
 
+  def test_make_array_from_single_device_arrays_tuple(self):
+    mesh = jtu.create_mesh((2, 2), ('x', 'y'))
+    shape = (8, 8)
+    s = jax.sharding.NamedSharding(mesh, P('x', 'y'))
+    inp_data = np.arange(math.prod(shape)).reshape(shape)
+
+    arrays = tuple(
+        jax.device_put(inp_data[index], d)
+        for d, index in s.addressable_devices_indices_map(shape).items())
+
+    jax.make_array_from_single_device_arrays(shape, s, arrays)  # doesn't crash
+
   def test_make_array_from_single_device_arrays_bad_inputs(self):
     x = jnp.arange(10)
     mesh = jtu.create_mesh((2,), ('x',))