drafting feature

Author: LarsKue

bayesflow/networks/deep_set/equivariant_layer.py

@@ -137,7 +137,10 @@ def call(self, input_set: Tensor, training: bool = False, **kwargs) -> Tensor:
         output_set = ops.concatenate([input_set, invariant_summary], axis=-1)
 
         # Pass through final equivariant transform + residual
-        output_set = input_set + self.equivariant_fc(output_set, training=training)
+        out_fc = self.equivariant_fc(output_set, training=training)
+        out_projected = self.out_fc_projector(out_fc)
+        output_set = input_set + out_projected
+        # output_set = input_set + self.equivariant_fc(output_set, training=training)
         if self.layer_norm is not None:
             output_set = self.layer_norm(output_set, training=training)
 

bayesflow/networks/deep_set/invariant_layer.py

@@ -106,8 +106,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.pooling_layer(set_summary, training=training)
         set_summary = self.outer_fc(set_summary, training=training)
+        set_summary = self.outer_projector(set_summary, training=training)
         return set_summary
 
     @sanitize_input_shape