specify explicit build functions for approximators

Author: vpratz

bayesflow/approximators/approximator.py

@@ -12,17 +12,15 @@
 
 class Approximator(BackendApproximator):
     def build(self, data_shapes: dict[str, tuple[int] | dict[str, dict]]) -> None:
-        mock_data = keras.tree.map_shape_structure(keras.ops.zeros, data_shapes)
-        self.build_from_data(mock_data)
+        raise NotImplementedError
 
     @classmethod
     def build_adapter(cls, **kwargs) -> Adapter:
         # implemented by each respective architecture
         raise NotImplementedError
 
     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
+        raise NotImplementedError
 
     @classmethod
     def build_dataset(

bayesflow/approximators/continuous_approximator.py

@@ -7,7 +7,14 @@
 from bayesflow.adapters import Adapter
 from bayesflow.networks import InferenceNetwork, SummaryNetwork
 from bayesflow.types import Tensor
-from bayesflow.utils import filter_kwargs, logging, split_arrays, squeeze_inner_estimates_dict, concatenate_valid
+from bayesflow.utils import (
+    filter_kwargs,
+    logging,
+    split_arrays,
+    squeeze_inner_estimates_dict,
+    concatenate_valid,
+    concatenate_valid_shapes,
+)
 from bayesflow.utils.serialization import serialize, deserialize, serializable
 
 from .approximator import Approximator
@@ -60,6 +67,28 @@ def __init__(
         else:
             self.standardize_layers = {var: Standardization(trainable=False) for var in self.standardize}
 
+    def build(self, data_shapes: dict[str, tuple[int] | dict[str, dict]]) -> None:
+        summary_outputs_shape = None
+        inference_conditions_shape = data_shapes.get("inference_conditions", None)
+        if self.summary_network is not None:
+            self.summary_network.build(data_shapes["summary_variables"])
+            summary_outputs_shape = self.summary_network.compute_output_shape(data_shapes["summary_variables"])
+        inference_conditions_shape = concatenate_valid_shapes(
+            [inference_conditions_shape, summary_outputs_shape], axis=-1
+        )
+        self.inference_network.build(data_shapes["inference_variables"], inference_conditions_shape)
+        if self.standardize == "all":
+            self.standardize = [
+                var
+                for var in ["inference_variables", "summary_variables", "inference_conditions"]
+                if var in data_shapes
+            ]
+
+            self.standardize_layers = {var: Standardization(trainable=False) for var in self.standardize}
+        for var, layer in self.standardize_layers.items():
+            layer.build(data_shapes[var])
+        self.built = True
+
     @classmethod
     def build_adapter(
         cls,
@@ -120,16 +149,7 @@ def compile(
         return super().compile(*args, **kwargs)
 
     def build_from_data(self, adapted_data: dict[str, any]):
-        if self.standardize == "all":
-            self.standardize = [
-                var
-                for var in ["inference_variables", "summary_variables", "inference_conditions"]
-                if var in adapted_data
-            ]
-
-            self.standardize_layers = {var: Standardization(trainable=False) for var in self.standardize}
-
-        super().build_from_data(adapted_data)
+        self.build(keras.tree.map_structure(keras.ops.shape, adapted_data))
 
     def compile_from_config(self, config):
         self.compile(**deserialize(config))

bayesflow/approximators/model_comparison_approximator.py

@@ -8,7 +8,7 @@
 from bayesflow.networks import SummaryNetwork
 from bayesflow.simulators import ModelComparisonSimulator, Simulator
 from bayesflow.types import Tensor
-from bayesflow.utils import filter_kwargs, logging, concatenate_valid
+from bayesflow.utils import filter_kwargs, logging, concatenate_valid, concatenate_valid_shapes
 from bayesflow.utils.serialization import serialize, deserialize, serializable
 
 from .approximator import Approximator
@@ -66,11 +66,27 @@ def __init__(
         else:
             self.standardize_layers = {var: Standardization(trainable=False) for var in self.standardize}
 
-    def build_from_data(self, adapted_data: dict[str, any]):
+    def build(self, data_shapes: dict[str, tuple[int] | dict[str, dict]]) -> None:
+        summary_outputs_shape = None
+        classifier_conditions_shape = data_shapes.get("classifier_conditions", None)
+        if self.summary_network is not None:
+            self.summary_network.build(data_shapes["summary_variables"])
+            summary_outputs_shape = self.summary_network.compute_output_shape(data_shapes["summary_variables"])
+        classifier_conditions_shape = concatenate_valid_shapes(
+            [classifier_conditions_shape, summary_outputs_shape], axis=-1
+        )
+        self.classifier_network.build(classifier_conditions_shape)
+        self.logits_projector.build(self.classifier_network.compute_output_shape(classifier_conditions_shape))
         if self.standardize == "all":
-            self.standardize = [var for var in ["summary_variables", "classifier_conditions"] if var in adapted_data]
+            self.standardize = [var for var in ["summary_variables", "classifier_conditions"] if var in data_shapes]
+
             self.standardize_layers = {var: Standardization(trainable=False) for var in self.standardize}
-        super().build_from_data(adapted_data)
+        for var, layer in self.standardize_layers.items():
+            layer.build(data_shapes[var])
+        self.built = True
+
+    def build_from_data(self, adapted_data: dict[str, any]):
+        self.build(keras.tree.map_structure(keras.ops.shape(adapted_data)))
 
     @classmethod
     def build_adapter(

bayesflow/utils/__init__.py

@@ -77,6 +77,7 @@
 
 from .tensor_utils import (
     concatenate_valid,
+    concatenate_valid_shapes,
     expand,
     expand_as,
     expand_to,

bayesflow/utils/tensor_utils.py

@@ -4,7 +4,7 @@
 import keras
 import numpy as np
 
-from bayesflow.types import Tensor
+from bayesflow.types import Tensor, Shape
 from . import logging
 
 T = TypeVar("T")
@@ -20,6 +20,17 @@ def concatenate_valid(tensors: Sequence[Tensor | None], axis: int = 0) -> Tensor
     return keras.ops.concatenate(tensors, axis=axis)
 
 
+def concatenate_valid_shapes(tensor_shapes: Sequence[Shape | None], axis: int = 0) -> Shape | None:
+    tensor_shapes = [s for s in tensor_shapes if s is not None]
+    if not tensor_shapes:
+        return None
+
+    output_shape = tensor_shapes[0]
+    for s in tensor_shapes[1:]:
+        output_shape[axis] += s[axis]
+    return output_shape
+
+
 def expand(x: Tensor, n: int, side: str):
     if n < 0:
         raise ValueError(f"Cannot expand {n} times.")

tests/test_approximators/test_build.py

@@ -0,0 +1,17 @@
+import keras
+from tests.utils import check_combination_simulator_adapter
+
+
+def test_build(approximator, simulator, batch_size, adapter):
+    check_combination_simulator_adapter(simulator, adapter)
+
+    num_batches = 4
+    data = simulator.sample((num_batches * batch_size,))
+
+    batch = adapter(data)
+    batch = keras.tree.map_structure(keras.ops.convert_to_tensor, batch)
+    batch_shapes = keras.tree.map_structure(keras.ops.shape, batch)
+    approximator.build(batch_shapes)
+    for layer in approximator.standardize_layers.values():
+        assert layer.built
+        assert layer.count == 0