diff --git a/bayesflow/experimental/cif/conditional_gaussian.py b/bayesflow/experimental/cif/conditional_gaussian.py index 435fc7f77..4e55dcdf4 100644 --- a/bayesflow/experimental/cif/conditional_gaussian.py +++ b/bayesflow/experimental/cif/conditional_gaussian.py @@ -33,8 +33,8 @@ def __init__(self, depth: int = 4, width: int = 128, activation: str = "swish", """ super().__init__(**keras_kwargs(kwargs)) - self.means = MLP(depth=depth, width=width, activation=activation) - self.stds = MLP(depth=depth, width=width, activation=activation) + self.means = MLP([width] * depth, activation=activation) + self.stds = MLP([width] * depth, activation=activation) self.output_projector = keras.layers.Dense(None) def build(self, input_shape: Shape) -> None: diff --git a/bayesflow/networks/mlp/mlp.py b/bayesflow/networks/mlp/mlp.py index 2522153ff..d35d35c4d 100644 --- a/bayesflow/networks/mlp/mlp.py +++ b/bayesflow/networks/mlp/mlp.py @@ -20,10 +20,8 @@ class MLP(keras.Layer): def __init__( self, + widths: Sequence[int] = (256, 256), *, - depth: int = None, - width: int = None, - widths: Sequence[int] = None, activation: str = "mish", kernel_initializer: str = "he_normal", residual: bool = False, @@ -46,15 +44,8 @@ def __init__( Parameters ---------- - depth : int, optional - Number of layers in the MLP when `widths` is not explicitly provided. Must be - used together with `width`. Default is 2. - width : int, optional - Number of units per layer when `widths` is not explicitly provided. Must be used - together with `depth`. Default is 256. widths : Sequence[int], optional - Explicitly defines the number of hidden units per layer. If provided, `depth` and - `width` should not be specified. Default is None. + Defines the number of hidden units per layer, as well as the number of layers to be used. activation : str, optional Activation function applied in the hidden layers, such as "mish". Default is "mish". kernel_initializer : str, optional @@ -76,17 +67,6 @@ def __init__( super().__init__(**keras_kwargs(kwargs)) - if widths is not None: - if depth is not None or width is not None: - raise ValueError("Either specify 'widths' or 'depth' and 'width', not both.") - else: - if depth is None or width is None: - # use the default - depth = 2 - width = 256 - - widths = [width] * depth - self.res_blocks = [] for width in widths: self.res_blocks.append( diff --git a/tests/test_networks/conftest.py b/tests/test_networks/conftest.py index 48cd79ed2..5d646d9ee 100644 --- a/tests/test_networks/conftest.py +++ b/tests/test_networks/conftest.py @@ -22,7 +22,7 @@ def flow_matching(): from bayesflow.networks import FlowMatching return FlowMatching( - subnet_kwargs={"widths": None, "width": 64, "depth": 2}, + subnet_kwargs={"widths": [64, 64]}, integrate_kwargs={"method": "rk45", "steps": 100}, )