implement requested changes and improve activation

Author: LarsKue

bayesflow/networks/deep_set/deep_set.py

@@ -30,7 +30,7 @@ def __init__(
         mlp_widths_invariant_inner: Sequence[int] = (64, 64),
         mlp_widths_invariant_outer: Sequence[int] = (64, 64),
         mlp_widths_invariant_last: Sequence[int] = (64, 64),
-        activation: str = "gelu",
+        activation: str = "silu",
         kernel_initializer: str = "he_normal",
         dropout: int | float | None = 0.05,
         spectral_normalization: bool = False,
@@ -72,7 +72,7 @@ def __init__(
         mlp_widths_invariant_last : Sequence[int], optional
             Widths of the MLP layers in the final invariant transformation. Default is (64, 64).
         activation : str, optional
-            Activation function used throughout the network, such as "gelu". Default is "gelu".
+            Activation function used throughout the network, such as "gelu". Default is "silu".
         kernel_initializer : str, optional
             Initialization strategy for kernel weights, such as "he_normal". Default is "he_normal".
         dropout : int, float, or None, optional

bayesflow/networks/deep_set/equivariant_layer.py

@@ -88,7 +88,7 @@ def __init__(
         # Fully connected net + residual connection for an equivariant transform applied to each set member
         self.input_projector = layers.Dense(mlp_widths_equivariant[-1])
         self.equivariant_fc = MLP(
-            mlp_widths_equivariant[:-1],
+            mlp_widths_equivariant,
             dropout=dropout,
             activation=activation,
             kernel_initializer=kernel_initializer,

bayesflow/networks/deep_set/invariant_layer.py

@@ -68,7 +68,7 @@ def __init__(
 
         # Inner fully connected net for sum decomposition: inner( pooling( inner(set) ) )
         self.inner_fc = MLP(
-            mlp_widths_inner[:-1],
+            mlp_widths_inner,
             dropout=dropout,
             activation=activation,
             kernel_initializer=kernel_initializer,
@@ -77,7 +77,7 @@ def __init__(
         self.inner_projector = keras.layers.Dense(mlp_widths_inner[-1], kernel_initializer=kernel_initializer)
 
         self.outer_fc = MLP(
-            mlp_widths_outer[:-1],
+            mlp_widths_outer,
             dropout=dropout,
             activation=activation,
             kernel_initializer=kernel_initializer,
@@ -108,10 +108,10 @@ def call(self, input_set: Tensor, training: bool = False, **kwargs) -> Tensor:
         """
 
         set_summary = self.inner_fc(input_set, training=training)
-        set_summary = self.inner_projector(set_summary, training=training)
+        set_summary = self.inner_projector(set_summary)
         set_summary = self.pooling_layer(set_summary, training=training)
         set_summary = self.outer_fc(set_summary, training=training)
-        set_summary = self.outer_projector(set_summary, training=training)
+        set_summary = self.outer_projector(set_summary)
         return set_summary
 
     @sanitize_input_shape