Repository: jax-ml/jax
Pull request #26838 (2025-02-28 05:53:18)

Plumb layout through the creation of IFRT Arrays (roll-forward with fix).
Plumb layout through the creation of IFRT Arrays (roll-forward with fix).

Reverts 7f9e7473cfe7e2b3c4eb43ce6df916b3159c1cff



Patch:
diff --git a/jax/_src/array.py b/jax/_src/array.py
index a8e3ec3d0..08cab98cf 100644
--- a/jax/_src/array.py
+++ b/jax/_src/array.py
@@ -44,7 +44,7 @@ from jax._src.sharding import Sharding
 from jax._src.sharding_impls import (
     PmapSharding, SingleDeviceSharding,
     device_replica_id_map, hashed_index, num_addressable_indices, local_to_global_shape)  # pyformat: disable
-from jax._src.typing import ArrayLike, DLDeviceType
+from jax._src.typing import ArrayLike, DLDeviceType, DTypeLike
 from jax._src.util import safe_zip, unzip3, use_cpp_class, use_cpp_method, cache
 import numpy as np
 
@@ -984,7 +984,8 @@ def make_array_from_process_local_data(
 
 
 def make_array_from_single_device_arrays(
-    shape: Shape, sharding: Sharding, arrays: Sequence[basearray.Array]
+    shape: Shape, sharding: Sharding, arrays: Sequence[basearray.Array], *,
+    dtype: DTypeLike | None = None,
 ) -> ArrayImpl:
   r"""Returns a ``jax.Array`` from a sequence of ``jax.Array``\s each on a single device.
       Every device in input ``sharding``\'s mesh must have an array in ``arrays``\s.
@@ -997,6 +998,8 @@ def make_array_from_single_device_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``.
+    dtype: The dtype of the output ``jax.Array``. If not provided, the dtype of the first array in
+      ``arrays`` is used. If ``arrays`` is empty, the ``dtype`` argument must be provided.
 
   Returns:
     A global ``jax.Array``, sharded as ``sharding``, with shape equal to ``shape``, and with per-device
@@ -1027,10 +1030,32 @@ def make_array_from_single_device_arrays(
   For cases where you have a local array and want to convert it to a global
   jax.Array, use ``jax.make_array_from_process_local_data``.
   """
+  if isinstance(arrays, Sequence):
+    if arrays:
+      if dtype is None:
+        dtype = arrays[0].dtype
+      else:
+        if arrays[0].dtype != dtype:
+          raise ValueError(
+              "If `dtype` is provided to "
+              "`jax.make_array_from_single_device_arrays`, it must match the "
+              f"dtype of the arrays in `arrays`. Got dtype={dtype} and arrays "
+              f"dtype={arrays[0].dtype}`.")
+    else:
+      if not config.enable_empty_arrays.value:
+        raise ValueError(
+            "An empty `arrays` list, to make a non-addressable Array, is "
+            "supported only if `jax.config.enable_empty_arrays` is set to True."
+        )
+      if dtype is None:
+        raise ValueError(
+            "If `arrays` is empty, `dtype` must be provided via the `dtype` "
+            "argument to `jax.make_array_from_single_device_arrays`.")
+
   # All input arrays should be committed. Checking it is expensive on
   # single-controller systems.
   aval = core.update_aval_with_sharding(
-      core.ShapedArray(shape, arrays[0].dtype, weak_type=False), sharding)
+      core.ShapedArray(shape, dtype, weak_type=False), sharding)
   if dtypes.issubdtype(aval.dtype, dtypes.extended):
     return aval.dtype._rules.make_sharded_array(aval, sharding, arrays,
                                                 committed=True)


Test patch:
diff --git a/tests/array_test.py b/tests/array_test.py
index dbebe5169..391638a9c 100644
--- a/tests/array_test.py
+++ b/tests/array_test.py
@@ -23,6 +23,7 @@ import numpy as np
 
 import jax
 import jax.numpy as jnp
+from jax._src import config
 from jax._src import core
 from jax._src import dispatch
 from jax._src import op_shardings
@@ -832,6 +833,28 @@ class JaxArrayTest(jtu.JaxTestCase):
     np.array(h_tensor)
     self.assertIsNone(h_tensor._npy_value)
 
+  @config.enable_empty_arrays(True)
+  def test_make_array_from_single_device_arrays_no_dtype_error(self):
+    mesh = jtu.create_mesh((4, 2), ('x', 'y'))
+    s = jax.sharding.NamedSharding(mesh, P('x', 'y'))
+    with self.assertRaisesRegex(
+        ValueError,
+        'If `arrays` is empty, `dtype` must be provided via the `dtype` '
+        'argument to `jax.make_array_from_single_device_arrays`.'):
+      jax.make_array_from_single_device_arrays((8, 2), s, [])
+
+  def test_make_array_from_single_device_arrays_bad_dtype_error(self):
+    s = jax.sharding.SingleDeviceSharding(jax.devices()[0])
+    shape = (8, 2)
+    np_inp = np.arange(math.prod(shape)).reshape(shape)
+    arr = jax.device_put(np_inp, s)
+    with self.assertRaisesRegex(
+        ValueError,
+        'If `dtype` is provided to `jax.make_array_from_single_device_arrays`, '
+        'it must match the dtype of the arrays in `arrays`.'):
+      jax.make_array_from_single_device_arrays(
+          shape, s, [arr], dtype=jnp.float32)
+
 
 class ShardingTest(jtu.JaxTestCase):