Merge pull request #306 from vpratz/fix-adapter-serialization

Fix: make transforms AsSet and AsTimeSeries serializable

Author: vpratz

bayesflow/adapters/transforms/as_set.py

@@ -33,3 +33,10 @@ def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
             return np.squeeze(data, axis=2)
 
         return data
+
+    @classmethod
+    def from_config(cls, config: dict, custom_objects=None) -> "AsSet":
+        return cls()
+
+    def get_config(self) -> dict:
+        return {}

bayesflow/adapters/transforms/as_time_series.py

@@ -30,3 +30,10 @@ def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
             return np.squeeze(data, axis=2)
 
         return data
+
+    @classmethod
+    def from_config(cls, config: dict, custom_objects=None) -> "AsTimeSeries":
+        return cls()
+
+    def get_config(self) -> dict:
+        return {}

bayesflow/networks/flow_matching/flow_matching.py

@@ -191,7 +191,11 @@ def compute_metrics(
         else:
             # not pre-configured, resample
             x1 = x
-            x0 = self.base_distribution.sample(keras.ops.shape(x1), seed=self.seed_generator)
+            if not self.built:
+                xz_shape = keras.ops.shape(x1)
+                conditions_shape = None if conditions is None else keras.ops.shape(conditions)
+                self.build(xz_shape, conditions_shape)
+            x0 = self.base_distribution.sample(keras.ops.shape(x1)[:-1])
 
             if self.use_optimal_transport:
                 x1, x0, conditions = optimal_transport(