Preserve tuples in TypeSpecs.

PiperOrigin-RevId: 690793837

Author: tensorflower-gardener

tf_keras/saving/legacy/saved_model/json_utils.py

@@ -210,7 +210,7 @@ def get_json_type(obj):
             return {
                 "class_name": "TypeSpec",
                 "type_spec": type_spec_name,
-                "serialized": obj._serialize(),
+                "serialized": _encode_tuple(obj._serialize()),
             }
         except ValueError:
             raise ValueError(

tf_keras/saving/legacy/saved_model/json_utils_test.py

@@ -16,6 +16,7 @@
 """Tests the JSON encoder and decoder."""
 
 import enum
+from typing import Mapping
 
 import tensorflow.compat.v2 as tf
 
@@ -24,6 +25,16 @@
 from tf_keras.testing_infra import test_utils
 
 
+class _ExtensionType(tf.experimental.ExtensionType):
+    """An ExtensionType with multiple tuples and mappings."""
+
+    __name__ = "tf_keras.json_utils.test._ExtensionType"
+
+    x: tf.Tensor
+    xy: tuple[tf.Tensor, tf.Tensor]
+    kv: Mapping[str, tf.Tensor]
+
+
 class JsonUtilsTest(test_combinations.TestCase):
     def test_encode_decode_tensor_shape(self):
         metadata = {
@@ -64,6 +75,17 @@ def test_encode_decode_type_spec(self):
         ):
             loaded = json_utils.decode(string)
 
+    def test_encode_decode_extensiontype_spec(self):
+        instance = _ExtensionType(
+            x=tf.constant(1),
+            xy=(tf.constant(2), tf.constant(True)),
+            kv={"a": tf.constant("foo"), "b": tf.constant("bar")},
+        )
+        spec = tf.type_spec_from_value(instance)
+        string = json_utils.Encoder().encode(spec)
+        loaded = json_utils.decode(string)
+        self.assertEqual(spec, loaded)
+
     def test_encode_decode_enum(self):
         class Enum(enum.Enum):
             CLASS_A = "a"