make metrics serializable

It seems that metrics do not store their state, I'm not sure yet if this
is intended behavior.

Author: vpratz

bayesflow/metrics/functional/maximum_mean_discrepancy.py

@@ -2,12 +2,20 @@
 
 from bayesflow.types import Tensor
 from bayesflow.utils import issue_url
+from bayesflow.utils.serialization import serializable
+
+from typing import Literal
 
 from .kernels import gaussian, inverse_multiquadratic
 
 
+@serializable("bayesflow.metrics")
 def maximum_mean_discrepancy(
-    x: Tensor, y: Tensor, kernel: str = "inverse_multiquadratic", unbiased: bool = False, **kwargs
+    x: Tensor,
+    y: Tensor,
+    kernel: Literal["inverse_multiquadratic", "gaussian"] = "inverse_multiquadratic",
+    unbiased: bool = False,
+    **kwargs,
 ) -> Tensor:
     """Computes a mixture of Gaussian radial basis functions (RBFs) between the samples of x and y.
 

bayesflow/metrics/functional/root_mean_squared_error.py

@@ -2,8 +2,10 @@
 from keras import ops
 
 from bayesflow.types import Tensor
+from bayesflow.utils.serialization import serializable
 
 
+@serializable("bayesflow.metrics")
 def root_mean_squared_error(x1: Tensor, x2: Tensor, normalize: bool = False, **kwargs) -> Tensor:
     """Computes the (normalized) root mean squared error between samples x1 and x2.
 

bayesflow/metrics/maximum_mean_discrepancy.py

@@ -1,8 +1,6 @@
-from functools import partial
-
 import keras
 
-from bayesflow.utils.serialization import serializable
+from bayesflow.utils.serialization import deserialize, serializable, serialize
 from .functional import maximum_mean_discrepancy
 
 
@@ -17,10 +15,22 @@ def __init__(
     ):
         super().__init__(name=name, **kwargs)
         self.mmd = self.add_variable(shape=(), initializer="zeros", name="mmd")
-        self.mmd_fn = partial(maximum_mean_discrepancy, kernel=kernel, unbiased=unbiased)
+        self.kernel = kernel
+        self.unbiased = unbiased
 
     def update_state(self, x, y):
-        self.mmd.assign(keras.ops.cast(self.mmd_fn(x, y), self.dtype))
+        self.mmd.assign(
+            keras.ops.cast(maximum_mean_discrepancy(x, y, kernel=self.kernel, unbiased=self.unbiased), self.dtype)
+        )
 
     def result(self):
         return self.mmd.value
+
+    def get_config(self):
+        base_config = super().get_config()
+        config = {"kernel": self.kernel, "unbiased": self.unbiased}
+        return base_config | serialize(config)
+
+    @classmethod
+    def from_config(cls, config, custom_objects=None):
+        return cls(**deserialize(config, custom_objects=custom_objects))

bayesflow/metrics/root_mean_squard_error.py

@@ -1,12 +1,20 @@
-from functools import partial
 import keras
 
-from bayesflow.utils.serialization import serializable
+from bayesflow.utils.serialization import deserialize, serializable
 from .functional import root_mean_squared_error
 
 
 @serializable("bayesflow.metrics")
 class RootMeanSquaredError(keras.metrics.MeanMetricWrapper):
     def __init__(self, name="root_mean_squared_error", dtype=None, **kwargs):
-        fn = partial(root_mean_squared_error, **kwargs)
-        super().__init__(fn, name=name, dtype=dtype)
+        super().__init__(root_mean_squared_error, name=name, dtype=dtype, **kwargs)
+
+    def get_config(self):
+        base_config = super().get_config()
+        # fn is fixed and passed directly in the constructor
+        base_config.pop("fn")
+        return base_config
+
+    @classmethod
+    def from_config(cls, config, custom_objects=None):
+        return cls(**deserialize(config, custom_objects=custom_objects))

tests/test_metrics/__init__.py

@@ -0,0 +1,20 @@
+import pytest
+
+
+@pytest.fixture()
+def root_mean_squared_error():
+    from bayesflow.metrics import RootMeanSquaredError
+
+    return RootMeanSquaredError(normalize=True, name="rmse", dtype="float64")
+
+
+@pytest.fixture()
+def maximum_mean_discrepancy():
+    from bayesflow.metrics import MaximumMeanDiscrepancy
+
+    return MaximumMeanDiscrepancy(name="mmd", kernel="gaussian", unbiased=True, dtype="float64")
+
+
+@pytest.fixture(params=["root_mean_squared_error", "maximum_mean_discrepancy"])
+def metric(request):
+    return request.getfixturevalue(request.param)

tests/test_metrics/conftest.py

@@ -0,0 +1,12 @@
+from bayesflow.utils.serialization import serialize, deserialize
+import keras
+
+
+def test_serialize_deserialize(metric, random_samples):
+    metric.update_state(keras.random.normal((2, 3)), keras.random.normal((2, 3)))
+
+    serialized = serialize(metric)
+    deserialized = deserialize(serialized)
+    reserialized = serialize(deserialized)
+
+    assert reserialized == serialized