Fix DTypePolicy serialization (#19387)

* Fix dtype policy serialization

* Update comments

keras/dtype_policies/dtype_policy.py

@@ -69,6 +69,10 @@ def __new__(cls, name):
             return FloatDTypePolicy(name)
         return super().__new__(cls)
 
+    def __getnewargs__(self):
+        # To support `copy`, `deepcopy` and `pickle`
+        return (self._name,)
+
     def __init__(self, name):
         self._name = name
         self._compute_dtype = backend.floatx()

keras/dtype_policies/dtype_policy_test.py

@@ -61,6 +61,32 @@ def test_get_config_from_config(self):
         new_policy = DTypePolicy.from_config(config)
         self.assertEqual(new_policy.name, "mixed_float16")
 
+    def test_deepcopy(self):
+        """Test builtin serialization methods."""
+        import copy
+        import pickle
+
+        # copy.deepcopy
+        policy = DTypePolicy("mixed_float16")
+        copied_policy = copy.deepcopy(policy)
+        self.assertEqual(
+            repr(copied_policy), '<FloatDTypePolicy "mixed_float16">'
+        )
+        # copy.copy
+        copied_policy = copy.copy(policy)
+        self.assertEqual(
+            repr(copied_policy), '<FloatDTypePolicy "mixed_float16">'
+        )
+        # pickle
+        temp_dir = self.get_temp_dir()
+        with open(f"{temp_dir}/policy.pickle", "wb") as f:
+            pickle.dump(policy, f)
+        with open(f"{temp_dir}/policy.pickle", "rb") as f:
+            copied_policy = pickle.load(f)
+        self.assertEqual(
+            repr(copied_policy), '<FloatDTypePolicy "mixed_float16">'
+        )
+
 
 class FloatDTypePolicyTest(test_case.TestCase):
     def test_initialization_valid_name(self):