From 5aaaa8a895fd6ae4dc1c2401f841abaa82928c89 Mon Sep 17 00:00:00 2001
From: Daniel Ng <dnlng@google.com>
Date: Fri, 26 Apr 2024 15:02:05 -0700
Subject: [PATCH] Move dtype settings out of  metadata field into the root of
 Tensorstore spec

Before, dtype used to be in the metadata field of tensorstore spec because of it was the legacy way to config the dtype.  This setting doesn't understand the "str" name, hence, there was special logic to translate bfloat for example.

This CL moves it out of the metadata field and put the dtype directly into the Tensorstore spec to eliminate special dtype translation logic.  This will also add support of other quantized types such as int4.

PiperOrigin-RevId: 628519426
---
 .../array_serialization/serialization.py           | 14 +++++---------
 .../array_serialization/serialization_test.py      | 14 ++++++++++----
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/jax/experimental/array_serialization/serialization.py b/jax/experimental/array_serialization/serialization.py
index e94ecaeaf853..a023d886c2a6 100644
--- a/jax/experimental/array_serialization/serialization.py
+++ b/jax/experimental/array_serialization/serialization.py
@@ -83,19 +83,11 @@ async def create_async_array_from_callback(
 
 
 def _get_metadata(arr):
-  if arr.dtype == jnp.bfloat16:
-    # Tensorstore uses 'bfloat16', not '<V2'.
-    dtype = 'bfloat16'
-  else:
-    dtype = np.dtype(arr.dtype).str
   local_shape = arr.addressable_data(0).shape
   return {
-      'compressor': {
-          'id': 'zstd'
-      },
+      'compressor': {'id': 'zstd'},
       'shape': arr.shape,
       'chunks': np.array(np.maximum(1, local_shape)),
-      'dtype': dtype,
   }
 
 
@@ -220,6 +212,10 @@ async def async_serialize(
   if not _spec_has_metadata(tensorstore_spec):
     tensorstore_spec['metadata'] = _get_metadata(arr_inp)
 
+  # Set dtype if it's not in spec
+  if 'dtype' not in tensorstore_spec:
+    tensorstore_spec['dtype'] = jnp.dtype(arr_inp.dtype).name
+
   # If primary_host is None, all hosts will checkpoint. This is used
   # for checkpointing to local filesystem.
   if primary_host is None or jax.process_index() == primary_host:
diff --git a/jax/experimental/array_serialization/serialization_test.py b/jax/experimental/array_serialization/serialization_test.py
index 62e74d3373f9..23264cd0a770 100644
--- a/jax/experimental/array_serialization/serialization_test.py
+++ b/jax/experimental/array_serialization/serialization_test.py
@@ -205,12 +205,15 @@ def cb3(_):
       self.assertArraysEqual(np.asarray(s.data), np.array([], dtype=np.float32))
     self.assertEqual(m3.dtype, np.float32)
 
-  def test_checkpointing_with_bigger_shape_jax_array(self):
+  @parameterized.product(input_dtype=[np.int32, jax.numpy.bfloat16])
+  def test_checkpointing_with_bigger_shape_jax_array(self, input_dtype):
     global_mesh = jtu.create_global_mesh((2, 2), ('x', 'y'))
     global_input_shape = (8, 2)
     num = math.prod(global_input_shape)
 
-    global_input_data1 = np.arange(num, dtype=np.int32).reshape(global_input_shape)
+    global_input_data1 = np.arange(num, dtype=input_dtype).reshape(
+        global_input_shape
+    )
     def cb1(index):
       return global_input_data1[index]
     arr = array.make_array_from_callback(
@@ -250,12 +253,15 @@ def cb1(index):
     for l in m2.addressable_shards:
       self.assertArraysEqual(l.data, global_input_data1.astype('float32'))
 
-  def test_checkpointing_with_int4(self):
+  @parameterized.product(input_dtype=[jax.numpy.int4, jax.numpy.int8])
+  def test_checkpointing_with_int4(self, input_dtype):
     global_mesh = jtu.create_global_mesh((2, 2), ('x', 'y'))
     global_input_shape = (8, 2)
     num = math.prod(global_input_shape)
 
-    global_input_data = np.arange(num, dtype=jax.numpy.int8).reshape(global_input_shape)
+    global_input_data = np.arange(num, dtype=input_dtype).reshape(
+        global_input_shape
+    )
     def cb(index):
       return global_input_data[index]
     arr = array.make_array_from_callback(