[WIP] Move standardization into approximators and make adapter stateless.

bayesflow/adapters/transforms/standardize.py

@@ -1,4 +1,5 @@
 from collections.abc import Sequence
+import warnings
 
 import numpy as np
 
@@ -69,6 +70,14 @@ def __init__(
     ):
         super().__init__()
 
+        if mean is None or std is None:
+            warnings.warn(
+                "Dynamic standardization is deprecated and will be removed in later versions."
+                "Instead, use the standardize argument of the approximator / workflow instance or provide "
+                "fixed mean and std arguments. You may incur some redundant computations if you keep this transform.",
+                DeprecationWarning,
+            )
+
         self.mean = mean
         self.std = std
 

bayesflow/adapters/transforms/transform.py

@@ -22,6 +22,7 @@ def __repr__(self):
 
     @classmethod
     def from_config(cls, config: dict, custom_objects=None):
+        # noinspection PyArgumentList
         return cls(**deserialize(config, custom_objects=custom_objects))
 
     def get_config(self) -> dict:

bayesflow/approximators/approximator.py

@@ -11,17 +11,17 @@
 
 
 class Approximator(BackendApproximator):
-    def build(self, data_shapes: any) -> None:
-        mock_data = keras.tree.map_structure(keras.ops.zeros, data_shapes)
+    def build(self, data_shapes: dict[str, tuple[int]]) -> None:
+        mock_data = keras.tree.map_shape_structure(keras.ops.zeros, data_shapes)
         self.build_from_data(mock_data)
 
     @classmethod
     def build_adapter(cls, **kwargs) -> Adapter:
         # implemented by each respective architecture
         raise NotImplementedError
 
-    def build_from_data(self, data: dict[str, any]) -> None:
-        self.compute_metrics(**filter_kwargs(data, self.compute_metrics), stage="training")
+    def build_from_data(self, adapted_data: dict[str, any]) -> None:
+        self.compute_metrics(**filter_kwargs(adapted_data, self.compute_metrics), stage="training")
         self.built = True
 
     @classmethod
@@ -61,6 +61,9 @@
             max_queue_size=max_queue_size,
         )
 
+    def call(self, *args, **kwargs):
+        return self.compute_metrics(*args, **kwargs)
+
     def fit(self, *, dataset: keras.utils.PyDataset = None, simulator: Simulator = None, **kwargs):
         """
         Trains the approximator on the provided dataset or on-demand data generated from the given simulator.
@@ -132,6 +135,7 @@
             logging.info("Building on a test batch.")
             mock_data = dataset[0]
             mock_data = keras.tree.map_structure(keras.ops.convert_to_tensor, mock_data)
-            self.build_from_data(mock_data)
+            mock_data_shapes = keras.tree.map_structure(keras.ops.shape, mock_data)
+            self.build(mock_data_shapes)
 
         return super().fit(dataset=dataset, **kwargs)