change build parameters to match Keras convention

The convention is to use parameter name with a `_shape` suffix.

Author: vpratz

bayesflow/networks/consistency_models/consistency_model.py

@@ -180,9 +180,9 @@ def build(self, xz_shape, conditions_shape=None):
         else:
             # Multiple separate inputs
             time_shape = tuple(xz_shape[:-1]) + (1,)  # same batch/sequence dims, 1 feature
-            self.subnet.build(input_shape_x=xz_shape, input_shape_t=time_shape, input_shape_conditions=conditions_shape)
+            self.subnet.build(x_shape=xz_shape, t_shape=time_shape, conditions_shape=conditions_shape)
             out_shape = self.subnet.compute_output_shape(
-                input_shape_x=xz_shape, input_shape_t=time_shape, input_shape_conditions=conditions_shape
+                x_shape=xz_shape, t_shape=time_shape, conditions_shape=conditions_shape
             )
         self.output_projector.build(out_shape)
 

bayesflow/networks/diffusion_model/diffusion_model.py

@@ -147,9 +147,9 @@ def build(self, xz_shape: Shape, conditions_shape: Shape = None) -> None:
         else:
             # Multiple separate inputs
             time_shape = tuple(xz_shape[:-1]) + (1,)  # same batch/sequence dims, 1 feature
-            self.subnet.build(input_shape_x=xz_shape, input_shape_t=time_shape, input_shape_conditions=conditions_shape)
+            self.subnet.build(x_shape=xz_shape, t_shape=time_shape, conditions_shape=conditions_shape)
             out_shape = self.subnet.compute_output_shape(
-                input_shape_x=xz_shape, input_shape_t=time_shape, input_shape_conditions=conditions_shape
+                x_shape=xz_shape, t_shape=time_shape, conditions_shape=conditions_shape
             )
 
         self.output_projector.build(out_shape)

bayesflow/networks/flow_matching/flow_matching.py

@@ -141,9 +141,9 @@ def build(self, xz_shape: Shape, conditions_shape: Shape = None) -> None:
         else:
             # Multiple separate inputs
             time_shape = tuple(xz_shape[:-1]) + (1,)  # same batch/sequence dims, 1 feature
-            self.subnet.build(input_shape_x=xz_shape, input_shape_t=time_shape, input_shape_conditions=conditions_shape)
+            self.subnet.build(x_shape=xz_shape, t_shape=time_shape, conditions_shape=conditions_shape)
             out_shape = self.subnet.compute_output_shape(
-                input_shape_x=xz_shape, input_shape_t=time_shape, input_shape_conditions=conditions_shape
+                x_shape=xz_shape, t_shape=time_shape, conditions_shape=conditions_shape
             )
 
         self.output_projector.build(out_shape)

tests/test_networks/conftest.py

@@ -18,7 +18,7 @@ def diffusion_model_edm_F():
     )
 
 
-@serializable("bayesflow.networks")
+@serializable("test", disable_module_check=True)
 class ConcatenateMLP(Sequential):
     def __init__(
         self,
@@ -33,27 +33,23 @@ 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_x, input_shape_t, input_shape_conditions=None):
+    def compute_output_shape(self, x_shape, t_shape, conditions_shape=None):
         concatenate_input_shapes = tuple(
             (
-                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),
+                x_shape[0],
+                x_shape[-1] + t_shape[-1] + (conditions_shape[-1] if conditions_shape is not None else 0),
             )
         )
         return self.mlp.compute_output_shape(concatenate_input_shapes)
 
-    def build(self, input_shape_x, input_shape_t, input_shape_conditions=None):
+    def build(self, x_shape, t_shape, conditions_shape=None):
         if self.built:
             return
 
         concatenate_input_shapes = tuple(
             (
-                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),
+                x_shape[0],
+                x_shape[-1] + t_shape[-1] + (conditions_shape[-1] if conditions_shape is not None else 0),
             )
         )
         self.mlp.build(concatenate_input_shapes)