deprecate (de)serialize_value_or_type

- replaced by serialize and deserialize
- adapt classes in the point inference network module which relied on the old
  structure
- add more lenient tests to compare configs
- adapt developer docs

Author: vpratz

bayesflow/links/ordered.py

@@ -1,8 +1,8 @@
 import keras
-from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.utils import layer_kwargs
 from bayesflow.utils.decorators import sanitize_input_shape
+from bayesflow.utils.serialization import deserialize, serializable, serialize
 
 
 @serializable(package="links.ordered")
@@ -14,12 +14,20 @@ def __init__(self, axis: int, anchor_index: int, **kwargs):
         self.axis = axis
         self.anchor_index = anchor_index
         self.group_indices = None
-
-        self.config = {"axis": axis, "anchor_index": anchor_index, **kwargs}
+        self._kwargs = kwargs
 
     def get_config(self):
         base_config = super().get_config()
-        return base_config | self.config
+        config = {
+            "axis": self.axis,
+            "anchor_index": self.anchor_index,
+            **self._kwargs,
+        }
+        return base_config | serialize(config)
+
+    @classmethod
+    def from_config(cls, config, custom_objects=None):
+        return cls(**deserialize(config, custom_objects=custom_objects))
 
     def build(self, input_shape):
         super().build(input_shape)

bayesflow/links/ordered_quantiles.py

@@ -1,7 +1,7 @@
 import keras
-from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.utils import layer_kwargs, logging
+from bayesflow.utils.serialization import serializable
 
 from collections.abc import Sequence
 
@@ -16,14 +16,12 @@ def __init__(self, q: Sequence[float] = None, axis: int = None, **kwargs):
         super().__init__(axis, None, **layer_kwargs(kwargs))
         self.q = q
 
-        self.config = {
-            "q": q,
-            "axis": axis,
-        }
-
     def get_config(self):
         base_config = super().get_config()
-        return base_config | self.config
+        config = {
+            "q": self.q,
+        }
+        return base_config | config
 
     def build(self, input_shape):
         if self.axis is None and 1 < len(input_shape) <= 3:

bayesflow/networks/point_inference_network.py

@@ -1,11 +1,7 @@
 import keras
-from keras.saving import (
-    deserialize_keras_object as deserialize,
-    serialize_keras_object as serialize,
-    register_keras_serializable as serializable,
-)
 
-from bayesflow.utils import model_kwargs, find_network, serialize_value_or_type, deserialize_value_or_type
+from bayesflow.utils import model_kwargs, find_network
+from bayesflow.utils.serialization import deserialize, serializable, serialize
 from bayesflow.types import Shape, Tensor
 from bayesflow.scores import ScoringRule, ParametricDistributionScore
 from bayesflow.utils.decorators import allow_batch_size
@@ -26,14 +22,9 @@ def __init__(
         super().__init__(**model_kwargs(kwargs))
 
         self.scores = scores
-
         self.subnet = find_network(subnet, **kwargs.get("subnet_kwargs", {}))
 
-        self.config = {
-            **kwargs,
-        }
-        self.config = serialize_value_or_type(self.config, "subnet", subnet)
-        self.config["scores"] = serialize(self.scores)
+        self._kwargs = kwargs
 
     def build(self, xz_shape: Shape, conditions_shape: Shape = None) -> None:
         """Builds all network components based on shapes of conditions and targets.
@@ -112,15 +103,18 @@ 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 | self.config
+        return base_config | serialize(config)
 
     @classmethod
     def from_config(cls, config):
         config = config.copy()
-        config["scores"] = deserialize(config["scores"])
-        config = deserialize_value_or_type(config, "subnet")
-        return cls(**config)
+        return cls(**deserialize(config))
 
     def call(
         self,

bayesflow/scores/quantile_score.py

@@ -1,11 +1,11 @@
 from typing import Sequence
 
 import keras
-from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Shape, Tensor
 from bayesflow.utils import logging, weighted_mean
 from bayesflow.links import OrderedQuantiles
+from bayesflow.utils.serialization import deserialize, serializable, serialize
 
 from .scoring_rule import ScoringRule
 
@@ -31,13 +31,17 @@ def __init__(self, q: Sequence[float] = None, links=None, **kwargs):
         self._q = keras.ops.convert_to_tensor(q, dtype="float32")
         self.links = links or {"value": OrderedQuantiles(q=q)}
 
-        self.config = {
-            "q": q,
-        }
-
     def get_config(self):
         base_config = super().get_config()
-        return base_config | self.config
+        config = {
+            "q": self.q,
+            "links": self.links,
+        }
+        return base_config | serialize(config)
+
+    @classmethod
+    def from_config(cls, config, custom_objects=None):
+        return cls(**deserialize(config, custom_objects=custom_objects))
 
     def get_head_shapes_from_target_shape(self, target_shape: Shape) -> dict[str, tuple]:
         # keras.saving.load_model sometimes passes target_shape as a list, so we force a conversion

bayesflow/scores/scoring_rule.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.utils import find_network, serialize_value_or_type, deserialize_value_or_type
+from bayesflow.utils import find_network
+from bayesflow.utils.serialization import deserialize, serializable, serialize
 
 
 @serializable(package="bayesflow.scores")
@@ -48,28 +48,17 @@ def __init__(
         self.subnets_kwargs = subnets_kwargs or {}
         self.links = links or {}
 
-        self.config = {"subnets_kwargs": self.subnets_kwargs}
-
     def get_config(self):
-        self.config["subnets"] = {
-            key: serialize_value_or_type({}, "subnet", subnet) for key, subnet in self.subnets.items()
+        config = {
+            "subnets": self.subnets,
+            "links": self.links,
         }
-        self.config["links"] = {key: serialize_value_or_type({}, "link", link) for key, link in self.links.items()}
 
-        return self.config
+        return serialize(config)
 
     @classmethod
     def from_config(cls, config):
-        config = config.copy()
-        config["subnets"] = {
-            key: deserialize_value_or_type(subnet_dict, "subnet")["subnet"]
-            for key, subnet_dict in config["subnets"].items()
-        }
-        config["links"] = {
-            key: deserialize_value_or_type(link_dict, "link")["link"] for key, link_dict in config["links"].items()
-        }
-
-        return cls(**config)
+        return cls(**deserialize(config))
 
     def get_head_shapes_from_target_shape(self, target_shape: Shape) -> dict[str, Shape]:
         """Request a dictionary of names and output shapes of required heads from the score."""

bayesflow/utils/serialization.py

@@ -5,6 +5,7 @@
 import keras
 import numpy as np
 import sys
+from warnings import warn
 
 # this import needs to be exactly like this to work with monkey patching
 from keras.saving import deserialize_keras_object
@@ -19,77 +20,21 @@
 
 
 def serialize_value_or_type(config, name, obj):
-    """Serialize an object that can be either a value or a type
-    and add it to a copy of the supplied dictionary.
-
-    Parameters
-    ----------
-    config  : dict
-        Dictionary to add the serialized object to. This function does not
-        modify the dictionary in place, but returns a modified copy.
-    name    : str
-        Name of the obj that should be stored. Required for later deserialization.
-    obj     : object or type
-        The object to serialize. If `obj` is of type `type`, we use
-        `keras.saving.get_registered_name` to obtain the registered type name.
-        If it is not a type, we try to serialize it as a Keras object.
-
-    Returns
-    -------
-    updated_config  : dict
-        Updated dictionary with a new key `"_bayesflow_<name>_type"` or
-        `"_bayesflow_<name>_val"`. The prefix is used to avoid name collisions,
-        the suffix indicates how the stored value has to be deserialized.
-
-    Notes
-    -----
-    We allow strings or `type` parameters at several places to instantiate objects
-    of a given type (e.g., `subnet` in `CouplingFlow`). As `type` objects cannot
-    be serialized, we have to distinguish the two cases for serialization and
-    deserialization. This function is a helper function to standardize and
-    simplify this.
-    """
-    updated_config = config.copy()
-    if isinstance(obj, type):
-        updated_config[f"{PREFIX}{name}_type"] = keras.saving.get_registered_name(obj)
-    else:
-        updated_config[f"{PREFIX}{name}_val"] = keras.saving.serialize_keras_object(obj)
-    return updated_config
+    """This function is deprecated."""
+    warn(
+        "This method is deprecated. It was replaced by bayesflow.utils.serialization.serialize.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
 
 
 def deserialize_value_or_type(config, name):
-    """Deserialize an object that can be either a value or a type and add
-    it to the supplied dictionary.
-
-    Parameters
-    ----------
-    config  : dict
-        Dictionary containing the object to deserialize. If a type was
-        serialized, it should contain the key `"_bayesflow_<name>_type"`.
-        If an object was serialized, it should contain the key
-        `"_bayesflow_<name>_val"`. In a copy of this dictionary,
-        the item will be replaced with the key `name`.
-    name    : str
-        Name of the object to deserialize.
-
-    Returns
-    -------
-    updated_config  : dict
-        Updated dictionary with a new key `name`, with a value that is either
-        a type or an object.
-
-    See Also
-    --------
-    serialize_value_or_type
-    """
-    updated_config = config.copy()
-    if f"{PREFIX}{name}_type" in config:
-        updated_config[name] = keras.saving.get_registered_object(config[f"{PREFIX}{name}_type"])
-        del updated_config[f"{PREFIX}{name}_type"]
-    elif f"{PREFIX}{name}_val" in config:
-        updated_config[name] = keras.saving.deserialize_keras_object(config[f"{PREFIX}{name}_val"])
-        del updated_config[f"{PREFIX}{name}_val"]
-    return updated_config
+    """This function is deprecated."""
+    warn(
+        "This method is deprecated. It was replaced by bayesflow.utils.serialization.deserialize.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
 
 
 def deserialize(config: dict, custom_objects=None, safe_mode=True, **kwargs):

docsrc/source/development/serialization.md

@@ -21,8 +21,6 @@ We mainly provide three convenience functions:
 - The {py:func}`~bayesflow.utils.serialization.serialize` function, which adds support for serializing classes.
 - Its counterpart {py:func}`~bayesflow.utils.serialization.deserialize`, adds support to deserialize classes.
 
-_Note: The `(de)serialize_value_or_type` functions are made obsolete by the functions given above and will probably be deprecated soon._
-
 ## Usage
 
 To use the adapted serialization functions, you have to use them in the `get_config` and `from_config` method. Please refer to existing classes in the library for usage examples.

tests/test_networks/test_point_inference_network/test_point_inference_network.py

@@ -3,7 +3,7 @@
     deserialize_keras_object as deserialize,
     serialize_keras_object as serialize,
 )
-from tests.utils import assert_layers_equal
+from tests.utils import assert_layers_equal, assert_configs_equal
 import pytest
 
 
@@ -71,8 +71,7 @@ def test_save_and_load_quantile(tmp_path, quantile_point_inference_network, rand
     keras.saving.save_model(net, tmp_path / "model.keras")
     loaded = keras.saving.load_model(tmp_path / "model.keras")
 
-    print(net.get_config())
-    assert net.get_config() == loaded.get_config()
+    assert_configs_equal(net.get_config(), loaded.get_config())
 
     assert_layers_equal(net, loaded)
 

tests/test_scores/test_scores.py

@@ -17,7 +17,6 @@ def test_score_output(scoring_rule, random_conditions):
 
     # Using random random_conditions also as targets for the purpose of this test.
     head_shapes = scoring_rule.get_head_shapes_from_target_shape(random_conditions.shape)
-    print(scoring_rule.get_config())
     estimates = {}
     for key, output_shape in head_shapes.items():
         link = scoring_rule.get_link(key)

tests/utils/assertions.py

@@ -35,3 +35,14 @@ def assert_layers_equal(layer1: keras.Layer, layer2: keras.Layer):
     # this is turned off for now, see https://github.com/bayesflow-org/bayesflow/issues/412
     msg = f"Layers {layer1.name} and {layer2.name} have a different name."
     # assert layer1.name == layer2.name, msg
+
+
+def assert_configs_equal(config1, config2):
+    """Asserts that two configs are equal.
+
+    Ignores whether lists or tuples were used, a difference that sometimes arises for
+    the `input_shape` in `build_config` entries
+    """
+    config1 = keras.tree.lists_to_tuples(config1)
+    config2 = keras.tree.lists_to_tuples(config2)
+    assert config1 == config2