Merge branch 'dev' into point-estimation

Author: han-ol

.github/workflows/docs.yaml renamed to .github/workflows/multiversion-docs.yaml

@@ -1,12 +1,14 @@
 
 # From https://github.com/eeholmes/readthedoc-test/blob/main/.github/workflows/docs_pages.yml
-name: docs
+name: multiversion-docs
 
-# execute this workflow automatically when we push to master
 on:
-  push:
-    branches:
-      - master
+  workflow_dispatch:
+# execute this workflow automatically when we push to master or dev
+#   push:
+#     branches:
+#       - master
+#       - dev
 
 jobs:
 
@@ -17,13 +19,15 @@ jobs:
       - name: Checkout main
         uses: actions/checkout@v3
         with:
-          path: master
+          path: dev
+          fetch-depth: 0
+          fetch-tags: true
 
-      - name: Checkout gh-pages
+      - name: Checkout gh-pages-dev
         uses: actions/checkout@v3
         with:
-          path: gh-pages
-          ref: gh-pages
+          path: gh-pages-dev
+          ref: gh-pages-dev
 
       - name: Set up Python
         uses: actions/setup-python@v4
@@ -33,17 +37,21 @@ jobs:
 
       - name: Install dependencies
         run: |
-          cd ./master
+          cd ./dev
           python -m pip install .[docs]
+      - name: Create local branches
+        run: |
+          cd ./dev
+          git branch master remotes/origin/master
       - name: Make the Sphinx docs
         run: |
-          cd ./master/docsrc
+          cd ./dev/docsrc
           make clean
           make github
       - name: Commit changes to docs
         run: |
-          cd ./gh-pages
-          cp -R ../master/docs/* ./
+          cd ./gh-pages-dev
+          cp -R ../dev/docs/* ./
           git config --local user.email ""
           git config --local user.name "github-actions"
           git add -A

.gitignore

@@ -5,6 +5,11 @@ __pycache__/
 projects/
 */bayesflow.egg-info
 docsrc/_build/
+docsrc/_build_polyversion
+docsrc/.bf_doc_gen_venv
+docsrc/source/api
+docsrc/source/_examples
+docsrc/source/contributing.md
 build
 docs/
 

CONTRIBUTING.md

@@ -100,7 +100,7 @@ Make sure to occasionally also run multi-backend tests for your OS using [tox](h
 tox --parallel auto
 ```
 
-See [tox.ini](tox.ini) for details on the environment configurations.
+See `tox.ini` for details on the environment configurations.
 Multi-OS tests will automatically be run once you create a pull request.
 
 Note that to be backend-agnostic, your code must not:
@@ -137,12 +137,24 @@ z = keras.ops.convert_to_numpy(x)
 ### 4. Document your changes
 
 The documentation uses [sphinx](https://www.sphinx-doc.org/) and relies on [numpy style docstrings](https://numpydoc.readthedocs.io/en/latest/format.html) in classes and functions.
-The overall *structure* of the documentation is manually designed. This also applies to the API documentation. This has two implications for you:
 
-1. If you add to existing submodules, the documentation will update automatically (given that you use proper numpy docstrings).
-2. If you add a new submodule or subpackage, you need to add a file to `docsrc/source/api` and a reference to the new module to the appropriate section of `docsrc/source/api/bayesflow.rst`.
+Run the following command to install all necessary packages for setting up documentation generation:
 
-You can re-build the documentation with
+```
+pip install .[docs]
+```
+
+The overall *structure* of the documentation is manually designed, but the API documentation is auto-generated.
+
+You can re-build the current documentation with
+
+```bash
+cd docsrc
+make clean && make dev
+# in case of issues, try `make clean-all`
+```
+
+We also provide a multi-version documentation. To generate it, run
 
 ```bash
 cd docsrc

bayesflow/adapters/adapter.py

@@ -122,12 +122,18 @@ def as_time_series(self, keys: str | Sequence[str]):
         return self
 
     def broadcast(
-        self, keys: str | Sequence[str], *, to: str, expand: str | int | tuple = "left", exclude: int | tuple = -1
+        self,
+        keys: str | Sequence[str],
+        *,
+        to: str,
+        expand: str | int | tuple = "left",
+        exclude: int | tuple = -1,
+        squeeze: int | tuple = None,
     ):
         if isinstance(keys, str):
             keys = [keys]
 
-        transform = Broadcast(keys, to=to, expand=expand, exclude=exclude)
+        transform = Broadcast(keys, to=to, expand=expand, exclude=exclude, squeeze=squeeze)
         self.transforms.append(transform)
         return self
 

bayesflow/adapters/transforms/broadcast.py

@@ -53,7 +53,15 @@ class Broadcast(Transform):
     It is recommended to precede this transform with a :class:`bayesflow.adapters.transforms.ToArray` transform.
     """
 
-    def __init__(self, keys: Sequence[str], *, to: str, expand: str | int | tuple = "left", exclude: int | tuple = -1):
+    def __init__(
+        self,
+        keys: Sequence[str],
+        *,
+        to: str,
+        expand: str | int | tuple = "left",
+        exclude: int | tuple = -1,
+        squeeze: int | tuple = None,
+    ):
         super().__init__()
         self.keys = keys
         self.to = to
@@ -67,6 +75,7 @@ def __init__(self, keys: Sequence[str], *, to: str, expand: str | int | tuple =
             exclude = (exclude,)
 
         self.exclude = exclude
+        self.squeeze = squeeze
 
     @classmethod
     def from_config(cls, config: dict, custom_objects=None) -> "Broadcast":
@@ -75,6 +84,7 @@ def from_config(cls, config: dict, custom_objects=None) -> "Broadcast":
             to=deserialize(config["to"], custom_objects),
             expand=deserialize(config["expand"], custom_objects),
             exclude=deserialize(config["exclude"], custom_objects),
+            squeeze=deserialize(config["squeeze"], custom_objects),
         )
 
     def get_config(self) -> dict:
@@ -83,6 +93,7 @@ def get_config(self) -> dict:
             "to": serialize(self.to),
             "expand": serialize(self.expand),
             "exclude": serialize(self.exclude),
+            "squeeze": serialize(self.squeeze),
         }
 
     # noinspection PyMethodOverriding
@@ -115,6 +126,9 @@ def forward(self, data: dict[str, np.ndarray], **kwargs) -> dict[str, np.ndarray
 
             data[k] = np.broadcast_to(data[k], new_shape)
 
+            if self.squeeze is not None:
+                data[k] = np.squeeze(data[k], axis=self.squeeze)
+
         return data
 
     # noinspection PyMethodOverriding

bayesflow/approximators/continuous_approximator.py

@@ -11,7 +11,7 @@
 from bayesflow.adapters import Adapter
 from bayesflow.networks import InferenceNetwork, SummaryNetwork
 from bayesflow.types import Tensor
-from bayesflow.utils import logging, split_arrays
+from bayesflow.utils import filter_kwargs, logging, split_arrays
 from .approximator import Approximator
 
 
@@ -141,7 +141,7 @@ def sample(
     ) -> dict[str, np.ndarray]:
         conditions = self.adapter(conditions, strict=False, stage="inference", **kwargs)
         conditions = keras.tree.map_structure(keras.ops.convert_to_tensor, conditions)
-        conditions = {"inference_variables": self._sample(num_samples=num_samples, **conditions)}
+        conditions = {"inference_variables": self._sample(num_samples=num_samples, **conditions, **kwargs)}
         conditions = keras.tree.map_structure(keras.ops.convert_to_numpy, conditions)
         conditions = self.adapter(conditions, inverse=True, strict=False, **kwargs)
 
@@ -154,6 +154,7 @@ def _sample(
         num_samples: int,
         inference_conditions: Tensor = None,
         summary_variables: Tensor = None,
+        **kwargs,
     ) -> Tensor:
         if self.summary_network is None:
             if summary_variables is not None:
@@ -162,7 +163,9 @@ def _sample(
             if summary_variables is None:
                 raise ValueError("Summary variables are required when a summary network is present.")
 
-            summary_outputs = self.summary_network(summary_variables)
+            summary_outputs = self.summary_network(
+                summary_variables, **filter_kwargs(kwargs, self.summary_network.call)
+            )
 
             if inference_conditions is None:
                 inference_conditions = summary_outputs
@@ -180,18 +183,26 @@ def _sample(
         else:
             batch_shape = (num_samples,)
 
-        return self.inference_network.sample(batch_shape, conditions=inference_conditions)
+        return self.inference_network.sample(
+            batch_shape,
+            conditions=inference_conditions,
+            **filter_kwargs(kwargs, self.inference_network.sample),
+        )
 
-    def log_prob(self, data: dict[str, np.ndarray]) -> np.ndarray:
-        data = self.adapter(data, strict=False, stage="inference")
+    def log_prob(self, data: dict[str, np.ndarray], **kwargs) -> np.ndarray:
+        data = self.adapter(data, strict=False, stage="inference", **kwargs)
         data = keras.tree.map_structure(keras.ops.convert_to_tensor, data)
-        log_prob = self._log_prob(**data)
+        log_prob = self._log_prob(**data, **kwargs)
         log_prob = keras.ops.convert_to_numpy(log_prob)
 
         return log_prob
 
     def _log_prob(
-        self, inference_variables: Tensor, inference_conditions: Tensor = None, summary_variables: Tensor = None
+        self,
+        inference_variables: Tensor,
+        inference_conditions: Tensor = None,
+        summary_variables: Tensor = None,
+        **kwargs,
     ) -> Tensor:
         if self.summary_network is None:
             if summary_variables is not None:
@@ -200,11 +211,17 @@ def _log_prob(
             if summary_variables is None:
                 raise ValueError("Summary variables are required when a summary network is present.")
 
-            summary_outputs = self.summary_network(summary_variables)
+            summary_outputs = self.summary_network(
+                summary_variables, **filter_kwargs(kwargs, self.summary_network.call)
+            )
 
             if inference_conditions is None:
                 inference_conditions = summary_outputs
             else:
                 inference_conditions = keras.ops.concatenate([inference_conditions, summary_outputs], axis=-1)
 
-        return self.inference_network.log_prob(inference_variables, conditions=inference_conditions)
+        return self.inference_network.log_prob(
+            inference_variables,
+            conditions=inference_conditions,
+            **filter_kwargs(kwargs, self.inference_network.log_prob),
+        )

bayesflow/networks/consistency_models/consistency_model.py

@@ -7,7 +7,7 @@
 import numpy as np
 
 from bayesflow.types import Tensor
-from bayesflow.utils import find_network, keras_kwargs
+from bayesflow.utils import find_network, keras_kwargs, serialize_value_or_type, deserialize_value_or_type
 
 
 from ..inference_network import InferenceNetwork
@@ -88,6 +88,27 @@ def __init__(
 
         self.seed_generator = keras.random.SeedGenerator()
 
+        # serialization: store all parameters necessary to call __init__
+        self.config = {
+            "total_steps": total_steps,
+            "max_time": max_time,
+            "sigma2": sigma2,
+            "eps": eps,
+            "s0": s0,
+            "s1": s1,
+            **kwargs,
+        }
+        self.config = serialize_value_or_type(self.config, "subnet", subnet)
+
+    def get_config(self):
+        base_config = super().get_config()
+        return base_config | self.config
+
+    @classmethod
+    def from_config(cls, config):
+        config = deserialize_value_or_type(config, "subnet")
+        return cls(**config)
+
     def _schedule_discretization(self, step) -> float:
         """Schedule function for adjusting the discretization level `N` during
         the course of training.

bayesflow/networks/consistency_models/continuous_consistency_model.py

@@ -7,7 +7,16 @@
 import numpy as np
 
 from bayesflow.types import Tensor
-from bayesflow.utils import jvp, concatenate, find_network, keras_kwargs, expand_right_as, expand_right_to
+from bayesflow.utils import (
+    jvp,
+    concatenate,
+    find_network,
+    keras_kwargs,
+    expand_right_as,
+    expand_right_to,
+    serialize_value_or_type,
+    deserialize_value_or_type,
+)
 
 
 from ..inference_network import InferenceNetwork
@@ -62,6 +71,22 @@ def __init__(
 
         self.seed_generator = keras.random.SeedGenerator()
 
+        # serialization: store all parameters necessary to call __init__
+        self.config = {
+            "sigma_data": sigma_data,
+            **kwargs,
+        }
+        self.config = serialize_value_or_type(self.config, "subnet", subnet)
+
+    def get_config(self):
+        base_config = super().get_config()
+        return base_config | self.config
+
+    @classmethod
+    def from_config(cls, config):
+        config = deserialize_value_or_type(config, "subnet")
+        return cls(**config)
+
     def _discretize_time(self, num_steps: int, rho: float = 3.5, **kwargs):
         t = np.linspace(0.0, np.pi / 2, num_steps)
         times = np.exp((t - np.pi / 2) * rho) * np.pi / 2

bayesflow/networks/coupling_flow/coupling_flow.py

@@ -2,7 +2,7 @@
 from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Tensor
-from bayesflow.utils import find_permutation, keras_kwargs
+from bayesflow.utils import find_permutation, keras_kwargs, serialize_value_or_type, deserialize_value_or_type
 
 from .actnorm import ActNorm
 from .couplings import DualCoupling
@@ -58,13 +58,33 @@ def __init__(
 
             self.invertible_layers.append(DualCoupling(subnet, transform, **kwargs.get("coupling_kwargs", {})))
 
+        # serialization: store all parameters necessary to call __init__
+        self.config = {
+            "depth": depth,
+            "transform": transform,
+            "permutation": permutation,
+            "use_actnorm": use_actnorm,
+            "base_distribution": base_distribution,
+            **kwargs,
+        }
+        self.config = serialize_value_or_type(self.config, "subnet", subnet)
+
     # noinspection PyMethodOverriding
     def build(self, xz_shape, conditions_shape=None):
         super().build(xz_shape)
 
         for layer in self.invertible_layers:
             layer.build(xz_shape=xz_shape, conditions_shape=conditions_shape)
 
+    def get_config(self):
+        base_config = super().get_config()
+        return base_config | self.config
+
+    @classmethod
+    def from_config(cls, config):
+        config = deserialize_value_or_type(config, "subnet")
+        return cls(**config)
+
     def _forward(
         self, x: Tensor, conditions: Tensor = None, density: bool = False, training: bool = False, **kwargs
     ) -> Tensor | tuple[Tensor, Tensor]: