allow and change default base distribution to None

bayesflow-org · Jul 16, 2024 · ab635f7 · ab635f7
1 parent b908c91
commit ab635f7
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/bayesflow/networks/summary_network.py b/bayesflow/networks/summary_network.py
@@ -6,9 +6,8 @@
 
 
 class SummaryNetwork(keras.Layer):
-    def __init__(self, base_distribution: str = "normal", **kwargs):
+    def __init__(self, base_distribution: str = None, **kwargs):
         super().__init__(**keras_kwargs(kwargs))
-
         self.base_distribution = find_distribution(base_distribution)
 
     def build(self, input_shape):

diff --git a/bayesflow/utils/dispatch/find_distribution.py b/bayesflow/utils/dispatch/find_distribution.py
@@ -13,12 +13,14 @@ def _(name: str, *args, **kwargs):
             from bayesflow.distributions import DiagonalNormal
 
             distribution = DiagonalNormal(*args, **kwargs)
+        case "none":
+            distribution = None
         case other:
             raise ValueError(f"Unsupported distribution name '{other}'.")
 
     return distribution
 
 
 @find_distribution.register
-def _(constructor: type, *args, **kwargs):
-    return constructor(*args, **kwargs)
+def _(none: None, *args, **kwargs):
+    return None