fix input shape

Author: arrjon

bayesflow/networks/consistency_models/consistency_model.py

@@ -169,18 +169,28 @@ def build(self, xz_shape, conditions_shape=None):
         input_shape = list(xz_shape)
 
         if self._concatenate_subnet_input:
-            # time vector
+            # construct time vector
             input_shape[-1] += 1
-
             if conditions_shape is not None:
                 input_shape[-1] += conditions_shape[-1]
+            input_shape = tuple(input_shape)
+
+            self.subnet.build(input_shape)
+            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
 
-        input_shape = tuple(input_shape)
+            # 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)
+            self.subnet.build(input_shape)  # Pass list of shapes
+            out_shape = self.subnet.compute_output_shape(input_shape)
 
-        input_shape = self.subnet.compute_output_shape(input_shape)
-        self.output_projector.build(input_shape)
+        self.output_projector.build(out_shape)
 
         # Choose coefficient according to [2] Section 3.3
         self.c_huber = 0.00054 * ops.sqrt(xz_shape[-1])

bayesflow/networks/diffusion_model/diffusion_model.py

@@ -140,11 +140,23 @@ def build(self, xz_shape: Shape, conditions_shape: Shape = None) -> None:
             input_shape[-1] += 1
             if conditions_shape is not None:
                 input_shape[-1] += conditions_shape[-1]
+            input_shape = tuple(input_shape)
 
-        input_shape = tuple(input_shape)
+            self.subnet.build(input_shape)
+            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)
-        out_shape = self.subnet.compute_output_shape(input_shape)
         self.output_projector.build(out_shape)
 
     def get_config(self):

bayesflow/networks/flow_matching/flow_matching.py

@@ -130,15 +130,28 @@ def build(self, xz_shape: Shape, conditions_shape: Shape = None) -> None:
 
         input_shape = list(xz_shape)
         if self._concatenate_subnet_input:
-            # account for concatenating the time and conditions
+            # construct time vector
             input_shape[-1] += 1
             if conditions_shape is not None:
                 input_shape[-1] += conditions_shape[-1]
-        input_shape = tuple(input_shape)
+            input_shape = tuple(input_shape)
 
-        self.subnet.build(input_shape)
-        input_shape = self.subnet.compute_output_shape(input_shape)
-        self.output_projector.build(input_shape)
+            self.subnet.build(input_shape)
+            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.output_projector.build(out_shape)
 
     @classmethod
     def from_config(cls, config, custom_objects=None):

tests/test_networks/conftest.py

@@ -27,6 +27,13 @@ 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):
+        concatenate_input_shapes = tuple(
+            (input_shape[0][0], sum([shape[-1] for shape in input_shape if shape is not None]))
+        )
+        out = self.mlp.compute_output_shape(concatenate_input_shapes)
+        return out
+
 
 @pytest.fixture()
 def diffusion_model_edm_F_subnet_concatenate():
@@ -198,9 +205,12 @@ def typical_point_inference_network_subnet():
         "affine_coupling_flow",
         "spline_coupling_flow",
         "flow_matching",
+        pytest.param("flow_matching_subnet_concatenate"),
         "free_form_flow",
         "consistency_model",
+        pytest.param("consistency_model_subnet_concatenate"),
         pytest.param("diffusion_model_edm_F"),
+        pytest.param("diffusion_model_edm_F_subnet_concatenate"),
         pytest.param("diffusion_model_edm_noise", marks=pytest.mark.slow),
         pytest.param("diffusion_model_cosine_velocity", marks=pytest.mark.slow),
         pytest.param("diffusion_model_cosine_F", marks=pytest.mark.slow),