Merge branch 'dev' into feat-diffusion-model

Author: vpratz

bayesflow/links/ordered.py

@@ -2,6 +2,7 @@
 from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.utils import layer_kwargs
+from bayesflow.utils.decorators import sanitize_input_shape
 
 
 @serializable(package="links.ordered")
@@ -49,5 +50,6 @@ def call(self, inputs):
         x = keras.ops.concatenate([below, anchor_input, above], self.axis)
         return x
 
+    @sanitize_input_shape
     def compute_output_shape(self, input_shape):
         return input_shape

bayesflow/networks/summary_network.py

@@ -21,6 +21,7 @@ def build(self, input_shape):
         if self.base_distribution is not None:
             self.base_distribution.build(keras.ops.shape(z))
 
+    @sanitize_input_shape
     def compute_output_shape(self, input_shape):
         return keras.ops.shape(self.call(keras.ops.zeros(input_shape)))
 

bayesflow/networks/transformers/mab.py

@@ -4,6 +4,7 @@
 from bayesflow.networks import MLP
 from bayesflow.types import Tensor
 from bayesflow.utils import layer_kwargs
+from bayesflow.utils.decorators import sanitize_input_shape
 from bayesflow.utils.serialization import serializable
 
 
@@ -122,8 +123,10 @@ def call(self, seq_x: Tensor, seq_y: Tensor, training: bool = False, **kwargs) -
         return out
 
     # noinspection PyMethodOverriding
+    @sanitize_input_shape
     def build(self, seq_x_shape, seq_y_shape):
         self.call(keras.ops.zeros(seq_x_shape), keras.ops.zeros(seq_y_shape))
 
+    @sanitize_input_shape
     def compute_output_shape(self, seq_x_shape, seq_y_shape):
         return keras.ops.shape(self.call(keras.ops.zeros(seq_x_shape), keras.ops.zeros(seq_y_shape)))

bayesflow/networks/transformers/pma.py

@@ -4,6 +4,7 @@
 from bayesflow.networks import MLP
 from bayesflow.types import Tensor
 from bayesflow.utils import layer_kwargs
+from bayesflow.utils.decorators import sanitize_input_shape
 from bayesflow.utils.serialization import serializable
 
 from .mab import MultiHeadAttentionBlock
@@ -125,5 +126,6 @@ def call(self, input_set: Tensor, training: bool = False, **kwargs) -> Tensor:
         summaries = self.mab(seed_tiled, set_x_transformed, training=training, **kwargs)
         return ops.reshape(summaries, (ops.shape(summaries)[0], -1))
 
+    @sanitize_input_shape
     def compute_output_shape(self, input_shape):
         return keras.ops.shape(self.call(keras.ops.zeros(input_shape)))

bayesflow/networks/transformers/sab.py

@@ -1,6 +1,7 @@
 import keras
 
 from bayesflow.types import Tensor
+from bayesflow.utils.decorators import sanitize_input_shape
 from bayesflow.utils.serialization import serializable
 
 from .mab import MultiHeadAttentionBlock
@@ -16,6 +17,7 @@ class SetAttentionBlock(MultiHeadAttentionBlock):
     """
 
     # noinspection PyMethodOverriding
+    @sanitize_input_shape
     def build(self, input_set_shape):
         self.call(keras.ops.zeros(input_set_shape))
 
@@ -42,5 +44,6 @@ def call(self, input_set: Tensor, training: bool = False, **kwargs) -> Tensor:
         return super().call(input_set, input_set, training=training, **kwargs)
 
     # noinspection PyMethodOverriding
+    @sanitize_input_shape
     def compute_output_shape(self, input_set_shape):
         return keras.ops.shape(self.call(keras.ops.zeros(input_set_shape)))

bayesflow/utils/decorators.py

@@ -114,7 +114,7 @@ def callback(x):
 
 
 def sanitize_input_shape(fn: Callable):
-    """Decorator to replace the first dimension in input_shape with a dummy batch size if it is None"""
+    """Decorator to replace the first dimension in ..._shape arguments with a dummy batch size if it is None"""
 
     # The Keras functional API passes input_shape = (None, second_dim, third_dim, ...), which
     # causes problems when constructions like self.call(keras.ops.zeros(input_shape)) are used
@@ -126,5 +126,8 @@ def callback(input_shape: Shape) -> Shape:
             return tuple(input_shape)
         return input_shape
 
-    fn = argument_callback("input_shape", callback)(fn)
+    args = inspect.getfullargspec(fn).args
+    for arg in args:
+        if arg.endswith("_shape"):
+            fn = argument_callback(arg, callback)(fn)
     return fn

examples/From_ABC_to_BayesFlow.ipynb

@@ -38,7 +38,10 @@
    "outputs": [],
    "source": [
     "import numpy as np\n",
-    "import matplotlib.pyplot as plt"
+    "import matplotlib.pyplot as plt\n",
+    "import tempfile\n",
+    "from pathlib import Path\n",
+    "import platform"
    ]
   },
   {
@@ -322,7 +325,9 @@
     ")\n",
     "\n",
     "# generate a temporary SQLite DB\n",
-    "abc_id = abc.new(\"sqlite:////tmp/mjp.db\", observations)"
+    "prefix = \"sqlite:///\" if platform.system() == \"Windows\" else \"sqlite:////\"\n",
+    "db_path = (Path(tempfile.gettempdir()).absolute() / \"mjp.db\").as_uri().replace(\"file:///\", prefix)\n",
+    "abc_id = abc.new(db_path, observations)"
    ]
   },
   {

examples/SIR_Posterior_Estimation.ipynb

@@ -19,7 +19,11 @@
    "source": [
     "import os\n",
     "# Set to your favorite backend\n",
-    "os.environ[\"KERAS_BACKEND\"] = \"jax\""
+    "if \"KERAS_BACKEND\" not in os.environ:\n",
+    "    # set this to \"torch\", \"tensorflow\", or \"jax\"\n",
+    "    os.environ[\"KERAS_BACKEND\"] = \"tensorflow\"\n",
+    "else:\n",
+    "    print(f\"Using '{os.environ['KERAS_BACKEND']}' backend\")"
    ]
   },
   {

examples/Two_Moons_Starter.ipynb

@@ -24,7 +24,11 @@
    "source": [
     "import os\n",
     "# Set to your favorite backend\n",
-    "os.environ[\"KERAS_BACKEND\"] = \"jax\""
+    "if \"KERAS_BACKEND\" not in os.environ:\n",
+    "    # set this to \"torch\", \"tensorflow\", or \"jax\"\n",
+    "    os.environ[\"KERAS_BACKEND\"] = \"tensorflow\"\n",
+    "else:\n",
+    "    print(f\"Using '{os.environ['KERAS_BACKEND']}' backend\")"
    ]
   },
   {

tests/test_networks/test_summary_networks.py

@@ -25,6 +25,28 @@ def test_build(automatic, summary_network, random_set):
     assert summary_network.variables, "Model has no variables."
 
 
+@pytest.mark.parametrize("automatic", [True, False])
+def test_build_functional_api(automatic, summary_network, random_set):
+    if summary_network is None:
+        pytest.skip(reason="Nothing to do, because there is no summary network.")
+
+    assert summary_network.built is False
+
+    inputs = keras.layers.Input(shape=keras.ops.shape(random_set)[1:])
+    outputs = summary_network(inputs)
+    model = keras.Model(inputs=inputs, outputs=outputs)
+
+    if automatic:
+        model(random_set)
+    else:
+        model.build(keras.ops.shape(random_set))
+
+    assert model.built is True
+
+    # check the model has variables
+    assert summary_network.variables, "Model has no variables."
+
+
 def test_variable_batch_size(summary_network, random_set):
     if summary_network is None:
         pytest.skip(reason="Nothing to do, because there is no summary network.")