[JAX] Fix a small bug if shardings is tuple.

# Details

`jax.tree.map` requests all its arguments to have the same data type.
From ```[None] * len(tensorstore_specs) if global_shapes is None else global_shapes```,
The data type is already decided to be a list. So if we pass `sharding` or `tspecs` as a tuple, it will fail.

Here we add an explicit conversion to a list for sharding and tspecs.

PiperOrigin-RevId: 707576866

Author: maxwillzq

jax/experimental/array_serialization/serialization.py

@@ -420,10 +420,9 @@ def run_deserialization(shardings: Sequence[sharding.Sharding | Layout],
   async def _run_deserializer():
     # Object should be created once per process.
     byte_limiter = _LimitInFlightBytes(concurrent_bytes)
-
     future_arrays = jax.tree_util.tree_map(
         partial(async_deserialize, byte_limiter=byte_limiter),
-        shardings, tensorstore_specs,
+        list(shardings), list(tensorstore_specs),
         [None] * len(tensorstore_specs) if global_shapes is None else global_shapes,
         [None] * len(tensorstore_specs) if dtypes is None else dtypes)
     return await asyncio.gather(*future_arrays)

jax/experimental/array_serialization/serialization_test.py

@@ -50,6 +50,40 @@ class CheckpointTest(jtu.JaxTestCase):
   def _on_commit_callback(self, temp_ckpt_dir, final_ckpt_dir):
     os.rename(temp_ckpt_dir, final_ckpt_dir)
 
+  def test_deserialize_on_array_list(self):
+    global_mesh = jtu.create_mesh((2, 4), ('x', 'y'))
+    inp_shape = (16, 64)
+    pspec = P('x', 'y')
+    sharding = NamedSharding(global_mesh, pspec)
+    inputs = []
+    lambda_fn = lambda idx: src[idx]
+    num_arrays = 5
+    for _ in range(num_arrays):
+      src = jax.random.normal(jax.random.key(0), inp_shape)
+      inp = array.make_array_from_callback(inp_shape, sharding, lambda_fn)
+      inputs.append(inp)
+    ckpt_dir = pathlib.Path(self.create_tempdir().full_path)
+    tspecs = [
+        serialization.get_tensorstore_spec(f'{ckpt_dir}/array_{i}')
+        for i in range(num_arrays)
+    ]
+    inputs = tuple(inputs)
+    tspecs = tuple(tspecs)
+    manager = serialization.GlobalAsyncCheckpointManager()
+    manager.serialize(
+        inputs,
+        tspecs,
+        on_commit_callback=partial(
+            self._on_commit_callback, ckpt_dir, ckpt_dir
+        ),
+    )
+    manager.wait_until_finished()
+    shardings = tuple([sharding] * num_arrays)
+    restored_arrays = manager.deserialize(shardings, tspecs)
+    self.assertLen(restored_arrays, num_arrays)
+    for inp, deserialized_array in zip(inputs, restored_arrays):
+      self.assertArraysEqual(deserialized_array, inp)
+
   @jtu.skip_on_devices('cpu')
   def test_memory_consumption(self):
     global_mesh = jtu.create_mesh((2, 4), ('x', 'y'))