improved naming

Author: arrjon

bayesflow/experimental/diffusion_model/diffusion_model.py

@@ -15,6 +15,7 @@
     integrate,
     integrate_stochastic,
     logging,
+    tensor_utils,
 )
 from bayesflow.utils.serialization import serialize, deserialize, serializable
 
@@ -115,7 +116,7 @@ def __init__(
         if subnet == "mlp":
             subnet_kwargs = DiffusionModel.MLP_DEFAULT_CONFIG | subnet_kwargs
         self.subnet = find_network(subnet, **subnet_kwargs)
-        self._subnet_concatenated_input = subnet_kwargs.get("concatenated_input", True)
+        self._concatenate_subnet_input = kwargs.get("concatenate_subnet_input", True)
 
         self.output_projector = keras.layers.Dense(units=None, bias_initializer="zeros", name="output_projector")
 
@@ -149,7 +150,8 @@ def get_config(self):
             "prediction_type": self._prediction_type,
             "loss_type": self._loss_type,
             "integrate_kwargs": self.integrate_kwargs,
-            "subnet_concatenated_input": self._subnet_concatenated_input,
+            "concatenate_subnet_input": self._concatenate_subnet_input,
+            # we do not need to store subnet_kwargs
         }
         return base_config | serialize(config)
 
@@ -198,7 +200,7 @@ def convert_prediction_to_x(
             return (z + sigma_t**2 * pred) / alpha_t
         raise ValueError(f"Unknown prediction type {self._prediction_type}.")
 
-    def subnet_input(
+    def _subnet_input(
         self, xz: Tensor, log_snr: Tensor, conditions: Tensor = None, training: bool = False
     ) -> Tensor | tuple[Tensor, Tensor, Tensor]:
         """
@@ -221,11 +223,8 @@ def subnet_input(
         Tensor
             The concatenated input tensor for the subnet or a tuple of tensors if concatenation is disabled.
         """
-        if self._subnet_concatenated_input:
-            if conditions is None:
-                xtc = keras.ops.concatenate([xz, log_snr], axis=-1)
-            else:
-                xtc = keras.ops.concatenate([xz, log_snr, conditions], axis=-1)
+        if self._concatenate_subnet_input:
+            xtc = tensor_utils.concatenate_valid([xz, log_snr, conditions], axis=-1)
             return self.subnet(xtc, training=training)
         else:
             return self.subnet(xz, log_snr, conditions, training=training)
@@ -271,7 +270,9 @@ def velocity(
         log_snr_t = ops.broadcast_to(log_snr_t, ops.shape(xz)[:-1] + (1,))
         alpha_t, sigma_t = self.noise_schedule.get_alpha_sigma(log_snr_t=log_snr_t)
 
-        subnet_out = self.subnet_input(xz, self._transform_log_snr(log_snr_t), conditions=conditions, training=training)
+        subnet_out = self._subnet_input(
+            xz, self._transform_log_snr(log_snr_t), conditions=conditions, training=training
+        )
         pred = self.output_projector(subnet_out, training=training)
 
         x_pred = self.convert_prediction_to_x(pred=pred, z=xz, alpha_t=alpha_t, sigma_t=sigma_t, log_snr_t=log_snr_t)
@@ -490,7 +491,7 @@ def compute_metrics(
         diffused_x = alpha_t * x + sigma_t * eps_t
 
         # calculate output of the network
-        subnet_out = self.subnet_input(
+        subnet_out = self._subnet_input(
             diffused_x, self._transform_log_snr(log_snr_t), conditions=conditions, training=training
         )
         pred = self.output_projector(subnet_out, training=training)

bayesflow/networks/consistency_models/consistency_model.py

@@ -4,7 +4,7 @@
 import numpy as np
 
 from bayesflow.types import Tensor
-from bayesflow.utils import find_network, layer_kwargs, weighted_mean
+from bayesflow.utils import find_network, layer_kwargs, weighted_mean, tensor_utils
 from bayesflow.utils.serialization import deserialize, serializable, serialize
 
 from ..inference_network import InferenceNetwork
@@ -77,7 +77,7 @@ def __init__(
         subnet_kwargs = subnet_kwargs or {}
         if subnet == "mlp":
             subnet_kwargs = ConsistencyModel.MLP_DEFAULT_CONFIG | subnet_kwargs
-        self._subnet_concatenated_input = subnet_kwargs.get("concatenated_input", True)
+        self._concatenate_subnet_input = kwargs.get("concatenate_subnet_input", True)
 
         self.subnet = find_network(subnet, **subnet_kwargs)
         self.output_projector = keras.layers.Dense(
@@ -120,7 +120,7 @@ def get_config(self):
             "eps": self.eps,
             "s0": self.s0,
             "s1": self.s1,
-            "subnet_concatenated_input": self._subnet_concatenated_input,
+            "concatenate_subnet_input": self._concatenate_subnet_input,
             # we do not need to store subnet_kwargs
         }
 
@@ -258,7 +258,7 @@ def _inverse(self, z: Tensor, conditions: Tensor = None, training: bool = False,
             x = self.consistency_function(x_n, t, conditions=conditions, training=training)
         return x
 
-    def subnet_input(
+    def _subnet_input(
         self, x: Tensor, t: Tensor, conditions: Tensor = None, training: bool = False
     ) -> Tensor | tuple[Tensor, Tensor, Tensor]:
         """
@@ -281,11 +281,8 @@ def subnet_input(
         Tensor
             The concatenated input tensor for the subnet or a tuple of tensors if concatenation is disabled.
         """
-        if self._subnet_concatenated_input:
-            if conditions is None:
-                xtc = keras.ops.concatenate([x, t], axis=-1)
-            else:
-                xtc = keras.ops.concatenate([x, t, conditions], axis=-1)
+        if self._concatenate_subnet_input:
+            xtc = tensor_utils.concatenate_valid([x, t, conditions], axis=-1)
             return self.subnet(xtc, training=training)
         else:
             return self.subnet(x, t, conditions, training=training)
@@ -305,7 +302,7 @@ def consistency_function(self, x: Tensor, t: Tensor, conditions: Tensor = None,
             Whether internal layers (e.g., dropout) should behave in train or inference mode.
         """
 
-        subnet_out = self.subnet_input(x, t, conditions, training=training)
+        subnet_out = self._subnet_input(x, t, conditions, training=training)
         f = self.output_projector(subnet_out)
 
         # Compute skip and out parts (vectorized, since self.sigma2 is of shape (1, input_dim)

bayesflow/networks/flow_matching/flow_matching.py

@@ -12,6 +12,7 @@
     layer_kwargs,
     optimal_transport,
     weighted_mean,
+    tensor_utils,
 )
 from bayesflow.utils.serialization import serialize, deserialize, serializable
 from ..inference_network import InferenceNetwork
@@ -107,7 +108,7 @@ def __init__(
         subnet_kwargs = subnet_kwargs or {}
         if subnet == "mlp":
             subnet_kwargs = FlowMatching.MLP_DEFAULT_CONFIG | subnet_kwargs
-        self._subnet_concatenated_input = subnet_kwargs.get("concatenated_input", True)
+        self._concatenate_subnet_input = kwargs.get("concatenate_subnet_input", True)
 
         self.subnet = find_network(subnet, **subnet_kwargs)
         self.output_projector = keras.layers.Dense(units=None, bias_initializer="zeros", name="output_projector")
@@ -148,13 +149,13 @@ def get_config(self):
             "loss_fn": self.loss_fn,
             "integrate_kwargs": self.integrate_kwargs,
             "optimal_transport_kwargs": self.optimal_transport_kwargs,
+            "concatenate_subnet_input": self._concatenate_subnet_input,
             # we do not need to store subnet_kwargs
-            "subnet_concatenated_input": self._subnet_concatenated_input,
         }
 
         return base_config | serialize(config)
 
-    def subnet_input(
+    def _subnet_input(
         self, x: Tensor, t: Tensor, conditions: Tensor = None, training: bool = False
     ) -> Tensor | tuple[Tensor, Tensor, Tensor]:
         """
@@ -177,11 +178,8 @@ def subnet_input(
         Tensor
             The concatenated input tensor for the subnet or a tuple of tensors if concatenation is disabled.
         """
-        if self._subnet_concatenated_input:
-            if conditions is None:
-                xtc = keras.ops.concatenate([x, t], axis=-1)
-            else:
-                xtc = keras.ops.concatenate([x, t, conditions], axis=-1)
+        if self._concatenate_subnet_input:
+            xtc = tensor_utils.concatenate_valid([x, t, conditions], axis=-1)
             return self.subnet(xtc, training=training)
         else:
             return self.subnet(x, t, conditions, training=training)
@@ -191,7 +189,7 @@ def velocity(self, xz: Tensor, time: float | Tensor, conditions: Tensor = None,
         time = expand_right_as(time, xz)
         time = keras.ops.broadcast_to(time, keras.ops.shape(xz)[:-1] + (1,))
 
-        subnet_out = self.subnet_input(xz, time, conditions, training=training)
+        subnet_out = self._subnet_input(xz, time, conditions, training=training)
         return self.output_projector(subnet_out, training=training)
 
     def _velocity_trace(