not working: first steps to achieve serialization with the new functions

Author: vpratz

bayesflow/links/ordered_quantiles.py

@@ -1,7 +1,7 @@
 import keras
 
 from bayesflow.utils import layer_kwargs, logging
-from bayesflow.utils.serialization import serializable
+from bayesflow.utils.serialization import serializable, serialize, deserialize
 
 from collections.abc import Sequence
 
@@ -21,7 +21,11 @@ def get_config(self):
         config = {
             "q": self.q,
         }
-        return base_config | config
+        return base_config | serialize(config)
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**deserialize(config))
 
     def build(self, input_shape):
         if self.axis is None and 1 < len(input_shape) <= 3:

bayesflow/networks/point_inference_network.py

@@ -22,9 +22,14 @@ def __init__(
         super().__init__(**model_kwargs(kwargs))
 
         self.scores = scores
+
         self.subnet = find_network(subnet, **kwargs.get("subnet_kwargs", {}))
 
-        self._kwargs = kwargs
+        self.config = {
+            "subnet": serialize(subnet),
+            "scores": serialize(scores),
+            **kwargs,
+        }
 
     def build(self, xz_shape: Shape, conditions_shape: Shape = None) -> None:
         """Builds all network components based on shapes of conditions and targets.
@@ -103,18 +108,14 @@ def build_from_config(self, config):
 
     def get_config(self):
         base_config = super().get_config()
-        config = {
-            "scores": self.scores,
-            "subnet": self.subnet,
-            **self._kwargs,
-        }
 
-        return base_config | serialize(config)
+        return base_config | self.config
 
     @classmethod
     def from_config(cls, config):
         config = config.copy()
-        return cls(**deserialize(config))
+        config = deserialize(config)
+        return cls(**config)
 
     def call(
         self,

bayesflow/scores/mean_score.py

@@ -1,4 +1,4 @@
-from keras.saving import register_keras_serializable as serializable
+from bayesflow.utils.serialization import serializable
 
 from .normed_difference_score import NormedDifferenceScore
 
@@ -12,4 +12,6 @@ class MeanScore(NormedDifferenceScore):
 
     def __init__(self, **kwargs):
         super().__init__(k=2, **kwargs)
-        self.config = {}
+
+    def get_config(self):
+        return {}

bayesflow/scores/median_score.py

@@ -1,4 +1,4 @@
-from keras.saving import register_keras_serializable as serializable
+from bayesflow.utils.serialization import serializable
 
 from .normed_difference_score import NormedDifferenceScore
 
@@ -12,4 +12,6 @@ class MedianScore(NormedDifferenceScore):
 
     def __init__(self, **kwargs):
         super().__init__(k=1, **kwargs)
-        self.config = {}
+
+    def get_config(self):
+        return {}

bayesflow/scores/multivariate_normal_score.py

@@ -1,10 +1,10 @@
 import math
 
 import keras
-from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Shape, Tensor
 from bayesflow.links import PositiveDefinite
+from bayesflow.utils.serialization import deserialize, serializable, serialize
 
 from .parametric_distribution_score import ParametricDistributionScore
 
@@ -32,11 +32,17 @@ def __init__(self, dim: int = None, links: dict = None, **kwargs):
         self.dim = dim
         self.links = links or {"covariance": PositiveDefinite()}
 
-        self.config = {"dim": dim}
-
     def get_config(self):
         base_config = super().get_config()
-        return base_config | self.config
+        config = {
+            "dim": self.dim,
+            "links": self.links,
+        }
+        return base_config | serialize(config)
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**deserialize(config))
 
     def get_head_shapes_from_target_shape(self, target_shape: Shape) -> dict[str, Shape]:
         self.dim = target_shape[-1]

bayesflow/scores/normed_difference_score.py

@@ -1,5 +1,5 @@
 import keras
-from keras.saving import register_keras_serializable as serializable
+from bayesflow.utils.serialization import deserialize, serializable
 
 from bayesflow.types import Shape, Tensor
 from bayesflow.utils import weighted_mean
@@ -20,8 +20,6 @@ def __init__(self, k: int, **kwargs):
         #: Exponent to absolute difference
         self.k = k
 
-        self.config = {"k": k}
-
     def get_head_shapes_from_target_shape(self, target_shape: Shape):
         # keras.saving.load_model sometimes passes target_shape as a list, so we force a conversion
         target_shape = tuple(target_shape)
@@ -61,4 +59,9 @@ def score(self, estimates: dict[str, Tensor], targets: Tensor, weights: Tensor =
 
     def get_config(self):
         base_config = super().get_config()
-        return base_config | self.config
+        config = dict(k=self.k)
+        return base_config | config
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**deserialize(config))