Fix types for subnets

Author: stefanradev93

bayesflow/experimental/diffusion_model/diffusion_model.py

@@ -51,7 +51,7 @@ class DiffusionModel(InferenceNetwork):
     def __init__(
         self,
         *,
-        subnet: str | type = "mlp",
+        subnet: str | type | keras.Layer = "mlp",
         noise_schedule: Literal["edm", "cosine"] | NoiseSchedule | type = "edm",
         prediction_type: Literal["velocity", "noise", "F", "x"] = "F",
         loss_type: Literal["velocity", "noise", "F"] = "noise",
@@ -69,9 +69,9 @@ def __init__(
 
         Parameters
         ----------
-        subnet : str or type, optional
-            Architecture for the transformation network. Can be "mlp" or a custom network class.
-            Default is "mlp".
+        subnet : str, type or keras.Layer, optional
+            Architecture for the transformation network. Can be "mlp", a custom network class, or
+            a Layer object, e.g., `bayesflow.networks.MLP(widths=[32, 32])`. Default is "mlp".
         noise_schedule : {'edm', 'cosine'} or NoiseSchedule or type, optional
             Noise schedule controlling the diffusion dynamics. Can be a string identifier,
             a schedule class, or a pre-initialized schedule instance. Default is "edm".

bayesflow/networks/coupling_flow/coupling_flow.py

@@ -40,7 +40,7 @@ class CouplingFlow(InferenceNetwork):
 
     def __init__(
         self,
-        subnet: str | type = "mlp",
+        subnet: str | type | keras.Layer = "mlp",
         depth: int = 6,
         transform: str = "affine",
         permutation: str | None = "random",
@@ -68,9 +68,9 @@ def __init__(
 
         Parameters
         ----------
-        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".
+        subnet : str, type, or keras.Layer, optional
+            Architecture for the transformation network. Can be "mlp", a custom network class, or
+            a Layer object, e.g., `bayesflow.networks.MLP(widths=[32, 32])`. Default is "mlp".
         depth : int, optional
             The number of invertible layers in the model. Default is 6.
         transform : str, optional

bayesflow/networks/flow_matching/flow_matching.py

@@ -51,7 +51,7 @@ class FlowMatching(InferenceNetwork):
 
     def __init__(
         self,
-        subnet: str | keras.Layer = "mlp",
+        subnet: str | type | keras.Layer = "mlp",
         base_distribution: str | Distribution = "normal",
         use_optimal_transport: bool = False,
         loss_fn: str | keras.Loss = "mse",
@@ -74,8 +74,8 @@ def __init__(
         Parameters
         ----------
         subnet : str or keras.Layer, optional
-            The architecture used for the transformation network. Can be "mlp" or a custom
-            callable network. Default is "mlp".
+            Architecture for the transformation network. Can be "mlp", a custom network class, or
+            a Layer object, e.g., `bayesflow.networks.MLP(widths=[32, 32])`. Default is "mlp".
         base_distribution : str, optional
             The base probability distribution from which samples are drawn, such as "normal".
             Default is "normal".