add get_config for network base classes

- get config has to be manually specified in the base classes, so that
  the config is stored even when to subclass overrides get_config
- to preserve the auto_config behavior, we have to use the
  `python_utils.default` decorator from, which marks them as default methods.
  This allows detecting if a subclass has overridden them. This is the
  same mechanism that Keras uses
- moved setting the `custom_metrics` parameter after the
  `super().__init__` calls, as the tracking is managed in setattr
- extended some tests to use metrics

Author: vpratz

bayesflow/experimental/diffusion_model/diffusion_model.py

@@ -141,6 +141,9 @@ def build(self, xz_shape: Shape, conditions_shape: Shape = None) -> None:
     def get_config(self):
         base_config = super().get_config()
 
+        # base distribution is passed manually to InferenceNetwork parent class, do not store it here
+        base_config.pop("base_distribution")
+
         config = {
             "subnet": self.subnet,
             "noise_schedule": self.noise_schedule,

bayesflow/experimental/free_form_flow/free_form_flow.py

@@ -8,7 +8,6 @@
     find_network,
     jacobian,
     jvp,
-    model_kwargs,
     vjp,
     weighted_mean,
 )
@@ -240,7 +239,6 @@ def from_config(cls, config, custom_objects=None):
 
     def get_config(self):
         base_config = super().get_config()
-        base_config = model_kwargs(base_config)
 
         config = {
             "beta": self.beta,

bayesflow/networks/consistency_models/consistency_model.py

@@ -110,6 +110,9 @@ def from_config(cls, config, custom_objects=None):
     def get_config(self):
         base_config = super().get_config()
 
+        # base distribution is passed manually to InferenceNetwork parent class, do not store it here
+        base_config.pop("base_distribution")
+
         config = {
             "total_steps": self.total_steps,
             "subnet": self.subnet,

bayesflow/networks/coupling_flow/coupling_flow.py

@@ -90,7 +90,7 @@ def __init__(
             Keyword arguments forwarded to the affine or spline transforms
             (e.g., bins for splines)
         **kwargs
-            Additional keyword arguments passed to `InvertibleLayer`.
+            Additional keyword arguments passed to `InferenceNetwork`.
 
         """
         super().__init__(base_distribution=base_distribution, **kwargs)

bayesflow/networks/inference_network.py

@@ -2,11 +2,12 @@
 from collections.abc import Sequence
 
 import keras
+from keras.src.utils import python_utils
 
 from bayesflow.types import Shape, Tensor
 from bayesflow.utils import layer_kwargs, find_distribution
 from bayesflow.utils.decorators import allow_batch_size
-from bayesflow.utils.serialization import serializable
+from bayesflow.utils.serialization import serializable, serialize
 
 
 @serializable("bayesflow.networks")
@@ -38,8 +39,8 @@ def __init__(
         **kwargs
             Additional keyword arguments forwarded to the `keras.Layer` constructor.
         """
-        self.custom_metrics = metrics
         super().__init__(**layer_kwargs(kwargs))
+        self.custom_metrics = metrics
         self.base_distribution = find_distribution(base_distribution)
 
     def build(self, xz_shape: Shape, conditions_shape: Shape = None) -> None:
@@ -104,3 +105,10 @@ def compute_metrics(
                 metrics[metric.name] = metric(samples, x)
 
         return metrics
+
+    @python_utils.default
+    def get_config(self):
+        base_config = super().get_config()
+
+        config = {"metrics": self.custom_metrics, "base_distribution": self.base_distribution}
+        return base_config | serialize(config)

bayesflow/networks/point_inference_network.py

@@ -1,4 +1,6 @@
+from collections.abc import Sequence
 import keras
+from keras.src.utils import python_utils
 
 from bayesflow.utils import model_kwargs, find_network
 from bayesflow.utils.serialization import deserialize, serializable, serialize
@@ -17,9 +19,12 @@ def __init__(
         self,
         scores: dict[str, ScoringRule],
         subnet: str | keras.Layer = "mlp",
+        *,
+        metrics: Sequence[keras.Metric] | None = None,
         **kwargs,
     ):
         super().__init__(**model_kwargs(kwargs))
+        self.custom_metrics = metrics
 
         self.scores = scores
 
@@ -28,6 +33,7 @@ def __init__(
         self.config = {
             "subnet": serialize(subnet),
             "scores": serialize(scores),
+            "metrics": serialize(metrics),
             **kwargs,
         }
 
@@ -106,6 +112,7 @@ def build_from_config(self, config):
             for head_key, head in self.heads[score_key].items():
                 head.name = config["heads"][score_key][head_key]
 
+    @python_utils.default
     def get_config(self):
         base_config = super().get_config()
 
@@ -114,9 +121,7 @@ def get_config(self):
     @classmethod
     def from_config(cls, config):
         config = config.copy()
-        config["scores"] = deserialize(config["scores"])
-        config["subnet"] = deserialize(config["subnet"])
-        return cls(**config)
+        return cls(**deserialize(config))
 
     def call(
         self,

bayesflow/networks/summary_network.py

@@ -1,11 +1,12 @@
 import keras
+from keras.src.utils import python_utils
 from collections.abc import Sequence
 
 from bayesflow.metrics.functional import maximum_mean_discrepancy
 from bayesflow.types import Tensor
 from bayesflow.utils import layer_kwargs, find_distribution
 from bayesflow.utils.decorators import sanitize_input_shape
-from bayesflow.utils.serialization import serializable
+from bayesflow.utils.serialization import serializable, serialize
 
 
 @serializable("bayesflow.networks")
@@ -44,8 +45,8 @@ def __init__(self, base_distribution: str = None, *, metrics: Sequence[keras.Met
         **kwargs
             Additional keyword arguments forwarded to the `keras.Layer` constructor.
         """
-        self.custom_metrics = metrics
         super().__init__(**layer_kwargs(kwargs))
+        self.custom_metrics = metrics
         self.base_distribution = find_distribution(base_distribution)
 
     @sanitize_input_shape
@@ -86,3 +87,10 @@ def compute_metrics(self, x: Tensor, stage: str = "training", **kwargs) -> dict[
                     metrics[metric.name] = metric(outputs, samples)
 
         return metrics
+
+    @python_utils.default
+    def get_config(self):
+        base_config = super().get_config()
+
+        config = {"metrics": self.custom_metrics, "base_distribution": self.base_distribution}
+        return base_config | serialize(config)

tests/test_approximators/conftest.py

@@ -40,6 +40,7 @@ def continuous_approximator(adapter, inference_network, summary_network):
 
 @pytest.fixture()
 def point_inference_network():
+    from bayesflow.metrics import RootMeanSquaredError
     from bayesflow.networks import PointInferenceNetwork
     from bayesflow.scores import NormedDifferenceScore, QuantileScore, MultivariateNormalScore
 
@@ -51,11 +52,13 @@ def point_inference_network():
         ),
         subnet="mlp",
         subnet_kwargs=dict(widths=(32, 32)),
+        metrics=[RootMeanSquaredError(name="rmse")],
     )
 
 
 @pytest.fixture()
 def point_inference_network_with_multiple_parametric_scores():
+    from bayesflow.metrics import RootMeanSquaredError
     from bayesflow.networks import PointInferenceNetwork
     from bayesflow.scores import MultivariateNormalScore
 
@@ -64,6 +67,7 @@ def point_inference_network_with_multiple_parametric_scores():
             mvn1=MultivariateNormalScore(),
             mvn2=MultivariateNormalScore(),
         ),
+        metrics=[RootMeanSquaredError(name="rmse")],
     )
 
 
@@ -181,9 +185,10 @@ def validation_dataset(batch_size, adapter, simulator):
 
 @pytest.fixture()
 def mean_std_summary_network():
+    from bayesflow.metrics import MaximumMeanDiscrepancy
     from tests.utils import MeanStdSummaryNetwork
 
-    return MeanStdSummaryNetwork()
+    return MeanStdSummaryNetwork(metrics=[MaximumMeanDiscrepancy("mmd")])
 
 
 @pytest.fixture(params=["continuous_approximator", "point_approximator", "model_comparison_approximator"])