add documentation

Author: arrjon

bayesflow/experimental/diffusion_model/diffusion_model.py

@@ -85,6 +85,12 @@ def __init__(
             Additional keyword arguments passed to the noise schedule constructor. Default is None.
         integrate_kwargs : dict[str, any], optional
             Configuration dictionary for integration during training or inference. Default is None.
+        concatenate_subnet_input: bool, optional
+            Flag for advanced users to control whether all inputs to the subnet should be concatenated
+            into a single vector or passed as separate arguments. If set to False, the subnet
+            must accept three separate inputs: 'x' (noisy parameters), 't' (log signal-to-noise ratio),
+            and optional 'conditions'. Default is True.
+
         **kwargs
             Additional keyword arguments passed to the base class and internal components.
         """
@@ -227,7 +233,7 @@ def _apply_subnet(
             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)
+            return self.subnet(x=xz, t=log_snr, conditions=conditions, training=training)
 
     def velocity(
         self,

bayesflow/networks/consistency_models/consistency_model.py

@@ -67,6 +67,11 @@ def __init__(
             Final number of discretization steps
         subnet_kwargs: dict[str, any], optional
             Keyword arguments passed to the subnet constructor or used to update the default MLP settings.
+        concatenate_subnet_input: bool, optional
+            Flag for advanced users to control whether all inputs to the subnet should be concatenated
+            into a single vector or passed as separate arguments. If set to False, the subnet
+            must accept three separate inputs: 'x' (noisy parameters), 't' (time),
+            and optional 'conditions'. Default is True.
         **kwargs    : dict, optional, default: {}
             Additional keyword arguments
         """
@@ -268,7 +273,7 @@ def _apply_subnet(
         Parameters
         ----------
         x : Tensor
-            The input tensor for the diffusion model, typically of shape (..., D), but can vary.
+            The parameter tensor, typically of shape (..., D), but can vary.
         t : Tensor
             The time tensor, typically of shape (..., 1).
         conditions : Tensor, optional
@@ -285,7 +290,7 @@ def _apply_subnet(
             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)
+            return self.subnet(x=x, t=t, conditions=conditions, training=training)
 
     def consistency_function(self, x: Tensor, t: Tensor, conditions: Tensor = None, training: bool = False) -> Tensor:
         """Compute consistency function.

bayesflow/networks/flow_matching/flow_matching.py

@@ -91,6 +91,11 @@ def __init__(
             Additional keyword arguments for configuring optimal transport. Default is None.
         subnet_kwargs: dict[str, any], optional, deprecated
             Keyword arguments passed to the subnet constructor or used to update the default MLP settings.
+        concatenate_subnet_input: bool, optional
+            Flag for advanced users to control whether all inputs to the subnet should be concatenated
+            into a single vector or passed as separate arguments. If set to False, the subnet
+            must accept three separate inputs: 'x' (noisy parameters), 't' (time),
+            and optional 'conditions'. Default is True.
         **kwargs
             Additional keyword arguments passed to the subnet and other components.
         """
@@ -165,7 +170,7 @@ def _apply_subnet(
         Parameters
         ----------
         x : Tensor
-            The input tensor for the diffusion model, typically of shape (..., D), but can vary.
+            The parameter tensor, typically of shape (..., D), but can vary.
         t : Tensor
             The time tensor, typically of shape (..., 1).
         conditions : Tensor, optional
@@ -185,7 +190,7 @@ def _apply_subnet(
         else:
             if training is False:
                 t = keras.ops.broadcast_to(t, keras.ops.shape(x)[:-1] + (1,))
-            return self.subnet(x, t, conditions, training=training)
+            return self.subnet(x=x, t=t, conditions=conditions, training=training)
 
     def velocity(self, xz: Tensor, time: float | Tensor, conditions: Tensor = None, training: bool = False) -> Tensor:
         time = keras.ops.convert_to_tensor(time, dtype=keras.ops.dtype(xz))