use monkey-patching to enable type deserialization

Author: LarsKue

bayesflow/utils/serialization.py

@@ -1,13 +1,21 @@
 from copy import copy
 
+import builtins
 import inspect
 import keras
+import numpy as np
 
+# this import needs to be exactly like this to work with monkey patching
+from keras.saving import deserialize_keras_object
+
+from .context_managers import monkey_patch
 from .decorators import allow_args
 
 
 PREFIX = "_bayesflow_"
 
+_type_prefix = "__bayesflow_type__"
+
 
 def serialize_value_or_type(config, name, obj):
     """Serialize an object that can be either a value or a type
@@ -83,10 +91,27 @@ def deserialize_value_or_type(config, name):
     return updated_config
 
 
-def deserialize(obj, custom_objects=None, module_objects=None):
-    if inspect.isclass(obj):
-        return keras.saving.get_registered_object(obj, custom_objects=custom_objects, module_objects=module_objects)
-    return keras.saving.deserialize_keras_object(obj, custom_objects=custom_objects, module_objects=module_objects)
+def deserialize(obj, custom_objects=None, safe_mode=True, **kwargs):
+    with monkey_patch(deserialize_keras_object, deserialize) as original_deserialize:
+        if isinstance(obj, str) and obj.startswith(_type_prefix):
+            # we marked this as a type during serialization
+            obj = obj[len(_type_prefix) :]
+            tp = keras.saving.get_registered_object(
+                obj, custom_objects=custom_objects, module_objects=builtins.__dict__ | np.__dict__
+            )
+            if tp is None:
+                raise ValueError(
+                    f"Could not deserialize type {obj!r}. Make sure it is registered with "
+                    f"`keras.saving.register_keras_serializable` or pass it in `custom_objects`."
+                )
+            return tp
+        if inspect.isclass(obj):
+            # add this base case since keras does not cover it
+            return obj
+
+        obj = original_deserialize(obj, custom_objects=custom_objects, safe_mode=safe_mode, **kwargs)
+
+        return obj
 
 
 @allow_args
@@ -107,7 +132,7 @@ def serializable(cls, package=None, name=None):
 def serialize(obj):
     if isinstance(obj, (tuple, list, dict)):
         return keras.tree.map_structure(serialize, obj)
+    elif inspect.isclass(obj):
+        return _type_prefix + keras.saving.get_registered_name(obj)
 
-    if inspect.isclass(obj):
-        return keras.saving.get_registered_name(obj)
     return keras.saving.serialize_keras_object(obj)