fix docstring

Author: arrjon

bayesflow/experimental/diffusion_model/diffusion_model.py

@@ -204,7 +204,7 @@ def _subnet_input(
         self, xz: Tensor, log_snr: Tensor, conditions: Tensor = None, training: bool = False
     ) -> Tensor | tuple[Tensor, Tensor, Tensor]:
         """
-        Prepares the input for the subnet either by concatenating the latent variable `xz`,
+        Prepares and passes the input to the subnet either by concatenating the latent variable `xz`,
         the signal-to-noise ratio `log_snr`, and optional conditions or by returning them separately.
 
         Parameters
@@ -221,7 +221,7 @@ def _subnet_input(
         Returns
         -------
         Tensor
-            The concatenated input tensor for the subnet or a tuple of tensors if concatenation is disabled.
+            The output tensor from the subnet.
         """
         if self._concatenate_subnet_input:
             xtc = tensor_utils.concatenate_valid([xz, log_snr, conditions], axis=-1)

bayesflow/networks/consistency_models/consistency_model.py

@@ -262,7 +262,7 @@ def _subnet_input(
         self, x: Tensor, t: Tensor, conditions: Tensor = None, training: bool = False
     ) -> Tensor | tuple[Tensor, Tensor, Tensor]:
         """
-        Prepares the input for the subnet either by concatenating the latent variable `x`,
+        Prepares and passes the input to the subnet either by concatenating the latent variable `x`,
         the time `t`, and optional conditions or by returning them separately.
 
         Parameters
@@ -279,7 +279,7 @@ def _subnet_input(
         Returns
         -------
         Tensor
-            The concatenated input tensor for the subnet or a tuple of tensors if concatenation is disabled.
+            The output tensor from the subnet.
         """
         if self._concatenate_subnet_input:
             xtc = tensor_utils.concatenate_valid([x, t, conditions], axis=-1)

bayesflow/networks/flow_matching/flow_matching.py

@@ -159,7 +159,7 @@ def _subnet_input(
         self, x: Tensor, t: Tensor, conditions: Tensor = None, training: bool = False
     ) -> Tensor | tuple[Tensor, Tensor, Tensor]:
         """
-        Prepares the input for the subnet either by concatenating the latent variable `x`,
+        Prepares and passes the input to the subnet either by concatenating the latent variable `x`,
         the time `t`, and optional conditions or by returning them separately.
 
         Parameters
@@ -176,7 +176,7 @@ def _subnet_input(
         Returns
         -------
         Tensor
-            The concatenated input tensor for the subnet or a tuple of tensors if concatenation is disabled.
+            The output tensor from the subnet.
         """
         if self._concatenate_subnet_input:
             t = keras.ops.broadcast_to(t, keras.ops.shape(x)[:-1] + (1,))