Repository: google/jax (python)

Issue:
jax.random.normal(key, 1000) error message is bad
Come on:

```
In [16]: jax.random.normal(jax.random.key(0), 1000)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[16], line 1
----> 1 jax.random.normal(jax.random.key(0), 1000)

File ~/packages/jax/jax/_src/random.py:710, in normal(key, shape, dtype)
    707   raise ValueError(f"dtype argument to `normal` must be a float or complex dtype, "
    708                    f"got {dtype}")
    709 dtype = dtypes.canonicalize_dtype(dtype)
--> 710 shape = core.as_named_shape(shape)
    711 return _normal(key, shape, dtype)

File ~/packages/jax/jax/_src/core.py:2332, in as_named_shape(shape)
   2330 if isinstance(shape, NamedShape):
   2331   return shape
-> 2332 return NamedShape(*shape)

TypeError: jax._src.core.NamedShape() argument after * must be an iterable, not int
```

Pull request description:
make core.as_named_shape treat `int` like `tuple[int]`
fixes #21343

Patch:
diff --git a/jax/_src/array.py b/jax/_src/array.py
index 796d313ae..0485bbe2c 100644
--- a/jax/_src/array.py
+++ b/jax/_src/array.py
@@ -32,6 +32,7 @@ from jax._src import dispatch
 from jax._src import dtypes
 from jax._src import errors
 from jax._src import profiler
+from jax._src import tree_util
 from jax._src import xla_bridge
 from jax._src.interpreters import mlir
 from jax._src.interpreters import pxla
@@ -1068,18 +1069,22 @@ def shard_device_array(x, devices, indices, sharding):
   aval = api_util.shaped_abstractify(x)
   return pxla.batched_device_put(aval, sharding, shards, devices)
 
+def _hashable_index(idx):
+  return tree_util.tree_map(
+      lambda x: (x.start, x.stop) if type(x) == slice else x, idx)
+
 
 def shard_sharded_device_array_slow_path(x, devices, indices, sharding):
   candidates = defaultdict(list)
   bufs = [buf.data for buf in x.addressable_shards]
   arr_indices = tuple(x.sharding.devices_indices_map(x.shape).values())
   for buf, idx in safe_zip(bufs, arr_indices):
-    candidates[hashed_index(idx)].append(buf)
+    candidates[_hashable_index(idx)].append(buf)
 
   bufs = []
   for idx, device in safe_zip(indices, devices):
     # Look up all buffers that contain the correct slice of the logical array.
-    candidates_list = candidates[hashed_index(idx)]
+    candidates_list = candidates[_hashable_index(idx)]
     if not candidates_list:
       # This array isn't sharded correctly. Reshard it via host roundtrip.
       # TODO(skye): more efficient reshard?
diff --git a/jax/_src/core.py b/jax/_src/core.py
index 052573781..75cb474e8 100644
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -2356,6 +2356,8 @@ def join_named_shapes(*named_shapes):
 
 # TODO: Make canonicalize_shape return named shapes?
 def as_named_shape(shape) -> NamedShape:
+  if isinstance(shape, int):
+    shape = (shape,)
   if isinstance(shape, NamedShape):
     return shape
   return NamedShape(*shape)
diff --git a/jax/_src/dispatch.py b/jax/_src/dispatch.py
index 6ed3e0924..57a70b3ff 100644
--- a/jax/_src/dispatch.py
+++ b/jax/_src/dispatch.py
@@ -16,7 +16,6 @@
 from __future__ import annotations
 
 import atexit
-import collections
 from collections.abc import Callable, Iterator, Sequence
 import contextlib
 import dataclasses
@@ -364,38 +363,13 @@ def _different_device_order_reshard(x, target_sharding):
         old_hlo_sharding.tile_assignment_devices(), permute_order
     )
     new_hlo_sharding = xc.HloSharding.from_proto(new_op_sharding)
-    # TODO(yashkatariya): Enable this when HloSharding conversion is fixed in
-    # XLA.
-    # assert (new_op_sharding.tile_assignment_dimensions
-    #         == new_hlo_sharding.tile_assignment_dimensions())
-    # assert (new_op_sharding.tile_assignment_devices
-    #         == new_hlo_sharding.tile_assignment_devices())
-
-  new_sharding = GSPMDSharding(
-      target_sharding._device_assignment, new_hlo_sharding,
-      memory_kind=target_sharding.memory_kind)
-
-  old_device_to_index_buffer = collections.defaultdict()
-  old_index_to_buffer = collections.defaultdict()
-  for s in x.addressable_shards:
-    old_index_to_buffer[array.hashed_index(s.index)] = s.data
-    old_device_to_index_buffer[s.device] = (s.index, s.data)
-
-  new_arrays = []
-  for new_d, new_index in new_sharding.addressable_devices_indices_map(x.shape).items():
-    old_index, old_buf = old_device_to_index_buffer[new_d]
-    if old_index == new_index:
-      assert array._get_device(old_buf) == new_d, (
-          array._get_device(old_buf), new_d)
-      new_arrays.append(old_buf)
-    else:
-      old_buf = old_index_to_buffer[array.hashed_index(new_index)]
-      new_arrays.append(
-          pxla.batched_device_put(old_buf.aval, SingleDeviceSharding(new_d),
-                                  [old_buf], [new_d]))
 
-  new_x = array.ArrayImpl(
-      x.aval, new_sharding, new_arrays, committed=True, _skip_checks=True)
+  new_x = array.make_array_from_single_device_arrays(
+      x.shape,
+      GSPMDSharding(target_sharding._device_assignment, new_hlo_sharding,
+                    memory_kind=target_sharding.memory_kind),
+      x._arrays,
+  )
 
   _orig_get_and_check_device_assignment = pxla._get_and_check_device_assignment.fn
   pxla._get_and_check_device_assignment.fn = partial(


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index bc37a3798..3b56a4f70 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -4241,35 +4241,6 @@ class ArrayPjitTest(jtu.JaxTestCase):
     self.assertArraysEqual(out, np_inp)
     self.assertEqual(out.sharding, s2)
 
-  @parameterized.named_parameters(
-      ("8_2", (8, 2)),
-      ("8_384", (8, 384)),
-  )
-  def test_device_put_efficient_reshard_complex_mesh(self, shape):
-    if jax.device_count() < 8:
-      self.skipTest('Requires >= 8 devices')
-
-    dev = jax.devices()
-    mesh1 = jax.sharding.Mesh(
-        np.asarray(dev).reshape([1, 2, 2, 2]),
-        ('replica', 'data', 'seq', 'model'))
-    mesh2 = jax.sharding.Mesh(
-        np.asarray(jax.devices())
-        .reshape([1, 1, 2, 2, 2, 1])
-        .swapaxes(2, 3)
-        .reshape([1, 1, 4, 2, 1]),
-        ('replica', 'data', 'seq', 'model_q', 'model_kv'))
-
-    np_inp = jnp.arange(math.prod(shape)).reshape(shape)
-    s1 = NamedSharding(mesh1, P('model'))
-    s2 = NamedSharding(mesh2, P())
-
-    x_s1 = jax.device_put(np_inp, s1)
-
-    out = jax.device_put(x_s1, s2)
-    self.assertArraysEqual(out, np_inp)
-    self.assertEqual(out.sharding, s2)
-
   def test_convert_element_type_sharding(self):
     mesh = jtu.create_global_mesh((2, 2), ('x', 'y'))
     s = NamedSharding(mesh, P('x', 'y'))
diff --git a/tests/random_test.py b/tests/random_test.py
index 2c45d60cc..80e8ea76f 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -1458,6 +1458,14 @@ class JnpWithKeyArrayTest(jtu.JaxTestCase):
     self.check_shape(func, keys, fill_value)
     self.check_against_reference(func, func, keys, fill_value)
 
+  def test_int_shape(self):
+    # It's not clear if we want to accept ints as the shape argument; the point
+    # of this test is not to check the API functionality but rather to ensure
+    # this doesn't fail in core.py like it used to.
+    @jax.jit
+    def f():
+      jax.random.normal(jax.random.key(0), 1000)
+    f()  # don't crash
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())