add compute_output_shape

Author: arrjon

bayesflow/networks/consistency_models/consistency_model.py

@@ -179,17 +179,11 @@ def build(self, xz_shape, conditions_shape=None):
             out_shape = self.subnet.compute_output_shape(input_shape)
         else:
             # Multiple separate inputs
-            main_input_shape = xz_shape
             time_shape = xz_shape[:-1] + (1,)  # same batch/sequence dims, 1 feature
-
-            # Build subnet with multiple input shapes
-            input_shape = [main_input_shape, time_shape]
-            if conditions_shape is not None:
-                input_shape.append(conditions_shape)
-
-            self.subnet.build(input_shape)  # Pass list of shapes
-            out_shape = self.subnet.compute_output_shape(input_shape)
-
+            self.subnet.build(input_shape_x=xz_shape, input_shape_t=time_shape, input_shape_conditions=conditions_shape)
+            out_shape = self.subnet.compute_output_shape(
+                input_shape_x=xz_shape, input_shape_t=time_shape, input_shape_conditions=conditions_shape
+            )
         self.output_projector.build(out_shape)
 
         # Choose coefficient according to [2] Section 3.3

bayesflow/networks/diffusion_model/diffusion_model.py

@@ -146,16 +146,11 @@ def build(self, xz_shape: Shape, conditions_shape: Shape = None) -> None:
             out_shape = self.subnet.compute_output_shape(input_shape)
         else:
             # Multiple separate inputs
-            main_input_shape = xz_shape
             time_shape = xz_shape[:-1] + (1,)  # same batch/sequence dims, 1 feature
-
-            # Build subnet with multiple input shapes
-            input_shape = [main_input_shape, time_shape]
-            if conditions_shape is not None:
-                input_shape.append(conditions_shape)
-
-            self.subnet.build(input_shape)  # Pass list of shapes
-            out_shape = self.subnet.compute_output_shape(input_shape)
+            self.subnet.build(input_shape_x=xz_shape, input_shape_t=time_shape, input_shape_conditions=conditions_shape)
+            out_shape = self.subnet.compute_output_shape(
+                input_shape_x=xz_shape, input_shape_t=time_shape, input_shape_conditions=conditions_shape
+            )
 
         self.output_projector.build(out_shape)
 

bayesflow/networks/flow_matching/flow_matching.py

@@ -140,16 +140,11 @@ def build(self, xz_shape: Shape, conditions_shape: Shape = None) -> None:
             out_shape = self.subnet.compute_output_shape(input_shape)
         else:
             # Multiple separate inputs
-            main_input_shape = xz_shape
             time_shape = xz_shape[:-1] + (1,)  # same batch/sequence dims, 1 feature
-
-            # Build subnet with multiple input shapes
-            input_shape = [main_input_shape, time_shape]
-            if conditions_shape is not None:
-                input_shape.append(conditions_shape)
-
-            self.subnet.build(input_shape)  # Pass list of shapes
-            out_shape = self.subnet.compute_output_shape(input_shape)
+            self.subnet.build(input_shape_x=xz_shape, input_shape_t=time_shape, input_shape_conditions=conditions_shape)
+            out_shape = self.subnet.compute_output_shape(
+                input_shape_x=xz_shape, input_shape_t=time_shape, input_shape_conditions=conditions_shape
+            )
 
         self.output_projector.build(out_shape)
 

tests/test_networks/conftest.py

@@ -28,13 +28,25 @@ def call(self, x, t, conditions=None, training=False):
         con = concatenate_valid([x, t, conditions], axis=-1)
         return self.mlp(con)
 
-    def compute_output_shape(self, input_shape):
+    def compute_output_shape(self, input_shape_x, input_shape_t, input_shape_conditions=None):
         concatenate_input_shapes = tuple(
-            (input_shape[0][0], sum([shape[-1] for shape in input_shape if shape is not None]))
+            (
+                input_shape_x[0],
+                input_shape_x[-1]
+                + input_shape_t[-1]
+                + (input_shape_conditions[-1] if input_shape_conditions is not None else 0),
+            )
         )
         out = self.mlp.compute_output_shape(concatenate_input_shapes)
         return out
 
+    def build(self, input_shape_x, input_shape_t, input_shape_conditions=None):
+        if self.built:
+            return
+
+        input_shape = self.compute_output_shape(input_shape_x, input_shape_t, input_shape_conditions)
+        self.mlp.build(input_shape)
+
 
 @pytest.fixture()
 def diffusion_model_edm_F_subnet_concatenate():