Some syntactic sugar fixes

Author: stefanradev93

bayesflow/networks/consistency_models/consistency_model.py

@@ -45,15 +45,16 @@ def __init__(
         eps: float = 0.001,
         s0: int | float = 10,
         s1: int | float = 50,
+        subnet_kwargs: dict[str, any] = None,
         **kwargs,
     ):
         """Creates an instance of a consistency model (CM) to be used for standalone consistency training (CT).
 
         Parameters
         ----------
         total_steps : int
-            The total number of training steps, can be calculate as
-            number of epochs * number of batches
+            The total number of training steps, must be calculated as number of epochs * number of batches
+            and cannot be inferred during construction time.
         subnet      : str or type, optional, default: "mlp"
             A neural network type for the consistency model, will be
             instantiated using subnet_kwargs.
@@ -67,21 +68,20 @@ def __init__(
             Initial number of discretization steps
         s1          : int or float, optional, default: 50
             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.
         **kwargs    : dict, optional, default: {}
             Additional keyword arguments
         """
         super().__init__(base_distribution="normal", **keras_kwargs(kwargs))
 
         self.total_steps = float(total_steps)
 
+        subnet_kwargs = subnet_kwargs or {}
         if subnet == "mlp":
-            subnet_kwargs = ConsistencyModel.MLP_DEFAULT_CONFIG.copy()
-            subnet_kwargs.update(kwargs.get("subnet_kwargs", {}))
-        else:
-            subnet_kwargs = kwargs.get("subnet_kwargs", {})
+            subnet_kwargs = ConsistencyModel.MLP_DEFAULT_CONFIG | subnet_kwargs
 
         self.student = find_network(subnet, **subnet_kwargs)
-
         self.student_projector = keras.layers.Dense(units=None, bias_initializer="zeros", kernel_initializer="zeros")
 
         self.sigma2 = ops.convert_to_tensor(sigma2)
@@ -108,6 +108,7 @@ def __init__(
             "eps": eps,
             "s0": s0,
             "s1": s1,
+            "subnet_kwargs": subnet_kwargs,
             **kwargs,
         }
         self.config = serialize_value_or_type(self.config, "subnet", subnet)

bayesflow/networks/coupling_flow/coupling_flow.py

@@ -36,8 +36,8 @@ class CouplingFlow(InferenceNetwork):
 
     def __init__(
         self,
-        depth: int = 6,
         subnet: str | type = "mlp",
+        depth: int = 6,
         transform: str = "affine",
         permutation: str | None = "random",
         use_actnorm: bool = True,
@@ -62,11 +62,11 @@ def __init__(
 
         Parameters
         ----------
-        depth : int, optional
-            The number of invertible layers in the model. Default is 6.
         subnet : str or type, optional
             The architecture type used within the coupling layers. Can be a string
             identifier like "mlp" or a callable type. Default is "mlp".
+        depth : int, optional
+            The number of invertible layers in the model. Default is 6.
         transform : str, optional
             The type of transformation used in the coupling layers, such as "affine".
             Default is "affine".

bayesflow/networks/coupling_flow/couplings/single_coupling.py

@@ -27,11 +27,9 @@ class SingleCoupling(InvertibleLayer):
     def __init__(self, subnet: str | type = "mlp", transform: str = "affine", **kwargs):
         super().__init__(**keras_kwargs(kwargs))
 
+        subnet_kwargs = kwargs.get("subnet_kwargs", {})
         if subnet == "mlp":
-            subnet_kwargs = SingleCoupling.MLP_DEFAULT_CONFIG.copy()
-            subnet_kwargs.update(kwargs.get("subnet_kwargs", {}))
-        else:
-            subnet_kwargs = kwargs.get("subnet_kwargs", {})
+            subnet_kwargs = SingleCoupling.MLP_DEFAULT_CONFIG | subnet_kwargs
 
         self.network = find_network(subnet, **subnet_kwargs)
         self.transform = find_transform(transform, **kwargs.get("transform_kwargs", {}))

bayesflow/networks/deep_set/equivariant_module.py

@@ -104,11 +104,19 @@ def call(self, input_set: Tensor, training: bool = False, **kwargs) -> Tensor:
 
         Parameters
         ----------
-        #TODO
+        input_set : Tensor
+            The input tensor representing a set, with shape
+            (batch_size, ..., set_size, input_dim).
+        training : bool, optional
+            A flag indicating whether the model is in training mode. Default is False.
+        **kwargs : dict
+            Additional keyword arguments for compatibility with other functions.
 
         Returns
         -------
-        #TODO
+        Tensor
+            The transformed output tensor with the same shape as `input_set`, where
+            each element is processed through the equivariant transformation.
         """
 
         input_set = self.input_projector(input_set)

bayesflow/networks/flow_matching/flow_matching.py

@@ -87,6 +87,8 @@ def __init__(
             Additional keyword arguments for the integration process. Default is None.
         optimal_transport_kwargs : dict[str, any], optional
             Additional keyword arguments for configuring optimal transport. Default is None.
+        subnet_kwargs: dict[str, any], optional
+            Keyword arguments passed to the subnet constructor or used to update the default MLP settings.
         **kwargs
             Additional keyword arguments passed to the subnet and other components.
         """
@@ -103,7 +105,6 @@ def __init__(
         self.seed_generator = keras.random.SeedGenerator()
 
         subnet_kwargs = subnet_kwargs or {}
-
         if subnet == "mlp":
             subnet_kwargs = FlowMatching.MLP_DEFAULT_CONFIG | subnet_kwargs
 
@@ -116,6 +117,7 @@ def __init__(
             "use_optimal_transport": self.use_optimal_transport,
             "optimal_transport_kwargs": self.optimal_transport_kwargs,
             "integrate_kwargs": self.integrate_kwargs,
+            "subnet_kwargs": subnet_kwargs,
             **kwargs,
         }
         self.config = serialize_value_or_type(self.config, "subnet", subnet)