Merge branch 'dev' into fork/vpratz/docs-migration-advice

LarsKue · LarsKue · commit 4cc500e423ee · 2025-04-25T14:17:10.000-04:00
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
@@ -73,8 +73,11 @@ jobs:
           pytest -x -m "not slow"
 
       - name: Run Slow Tests
-        # run all slow tests only on manual trigger
-        if: github.event_name == 'workflow_dispatch'
+        # Run slow tests on manual trigger and pushes/PRs to main.
+        # Limit to one OS and Python version to save compute.
+        # Multiline if statements are weird, https://github.com/orgs/community/discussions/25641,
+        # but feel free to convert it.
+        if: ${{ ((github.event_name == 'workflow_dispatch') || (github.event_name == 'push' && github.ref_name == 'main') || (github.event_name == 'pull_request' && github.base_ref == 'main')) && ((matrix.os == 'windows-latest') && (matrix.python-version == '3.10')) }}
         run: |
           pytest -m "slow"
 
diff --git a/bayesflow/links/ordered.py b/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
diff --git a/bayesflow/networks/coupling_flow/coupling_flow.py b/bayesflow/networks/coupling_flow/coupling_flow.py
@@ -77,7 +77,7 @@ def __init__(
             The type of transformation used in the coupling layers, such as "affine".
             Default is "affine".
         permutation : str or None, optional
-            The type of permutation applied between layers. Can be "random" or None
+            The type of permutation applied between layers. Can be "orthogonal", "random", "swap", or None
             (no permutation). Default is "random".
         use_actnorm : bool, optional
             Whether to apply ActNorm before each coupling layer. Default is True.
diff --git a/bayesflow/networks/summary_network.py b/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)))
 
diff --git a/bayesflow/networks/transformers/mab.py b/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)))
diff --git a/bayesflow/networks/transformers/pma.py b/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)))
diff --git a/bayesflow/networks/transformers/sab.py b/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)))
diff --git a/bayesflow/utils/decorators.py b/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
diff --git a/docsrc/source/about.rst b/docsrc/source/about.rst
@@ -1,4 +1,4 @@
-About us
+About Us
 ========
 
 Core maintainers
diff --git a/docsrc/source/conf.py b/docsrc/source/conf.py
@@ -141,7 +141,29 @@
         "image_light": "_static/bayesflow_hor.png",
         "image_dark": "_static/bayesflow_hor_dark.png",
     },
-    "navbar_center": ["version-switcher", "navbar-nav"],
+    "icon_links_label": "Icon Links",
+    "icon_links": [
+        {
+            "name": "GitHub",
+            "url": "https://github.com/bayesflow-org/bayesflow",
+            "icon": "fa-brands fa-square-github",
+            "type": "fontawesome",
+        },
+        {
+            "name": "Discourse Forum",
+            "url": "https://discuss.bayesflow.org/",
+            "icon": "fa-brands fa-discourse",
+            "type": "fontawesome",
+        },
+    ],
+    "navbar_align": "left",
+    # -- Template placement in theme layouts ----------------------------------
+    "navbar_start": ["navbar-logo"],
+    # Note that the alignment of navbar_center is controlled by navbar_align
+    "navbar_center": ["navbar-nav"],
+    "navbar_end": ["theme-switcher", "navbar-icon-links", "version-switcher"],
+    # navbar_persistent is persistent right (even when on mobiles)
+    "navbar_persistent": ["search-button"],
     "switcher": {
         "json_url": "/versions.json",
         "version_match": current,
diff --git a/docsrc/source/index.md b/docsrc/source/index.md
@@ -240,6 +240,8 @@ If you are interested in a curated list of resources, including reviews, softwar
 
 This project is currently managed by researchers from Rensselaer Polytechnic Institute, TU Dortmund University, and Heidelberg University. It is partially funded by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation, Project 528702768). The project is further supported by Germany's Excellence Strategy -- EXC-2075 - 390740016 (Stuttgart Cluster of Excellence SimTech) and EXC-2181 - 390900948 (Heidelberg Cluster of Excellence STRUCTURES), as well as the Informatics for Life initiative funded by the Klaus Tschira Foundation.
 
+The [scikit-learn](https://scikit-learn.org/) website was a great resource and inspration for this site and the API documentation. We thank the scikit-learn community for sharing their configurations, which allowed us to include many nice features into this site as well.
+
 ## License \& Source Code
 
 BayesFlow is released under {mainbranch}`MIT License <LICENSE>`.
@@ -261,5 +263,5 @@ examples
 api/bayesflow
 about
 Contributing <contributing>
-Developer docs <development/index>
+Developer Docs <development/index>
 ```
diff --git a/examples/From_ABC_to_BayesFlow.ipynb b/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)"
    ]
   },
   {
diff --git a/examples/SIR_Posterior_Estimation.ipynb b/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\")"
    ]
   },
   {
diff --git a/examples/Two_Moons_Starter.ipynb b/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\")"
    ]
   },
   {
diff --git a/tests/test_networks/test_coupling_flow/test_permutations.py b/tests/test_networks/test_coupling_flow/test_permutations.py
@@ -0,0 +1,117 @@
+import pytest
+import keras
+import numpy as np
+
+from bayesflow.networks.coupling_flow.permutations import (
+    FixedPermutation,
+    OrthogonalPermutation,
+    RandomPermutation,
+    Swap,
+)
+
+
+@pytest.fixture(params=[FixedPermutation, OrthogonalPermutation, RandomPermutation, Swap])
+def permutation_class(request):
+    return request.param
+
+
+@pytest.fixture
+def input_tensor():
+    return keras.random.normal((2, 5))
+
+
+def test_fixed_permutation_build_and_call():
+    # Since FixedPermutation is abstract, create a subclass for testing build.
+    class TestPerm(FixedPermutation):
+        def build(self, xz_shape, **kwargs):
+            length = xz_shape[-1]
+            self.forward_indices = keras.ops.arange(length - 1, -1, -1)
+            self.inverse_indices = keras.ops.arange(length - 1, -1, -1)
+
+    layer = TestPerm()
+    input_shape = (2, 4)
+    layer.build(input_shape)
+
+    x = keras.ops.convert_to_tensor(np.arange(8).reshape(input_shape).astype("float32"))
+    z, log_det = layer(x, inverse=False)
+    x_inv, log_det_inv = layer(z, inverse=True)
+
+    # Check shape preservation
+    assert z.shape == x.shape
+    assert x_inv.shape == x.shape
+    # Forward then inverse recovers input
+    np.testing.assert_allclose(keras.ops.convert_to_numpy(x_inv), keras.ops.convert_to_numpy(x), atol=1e-5)
+    # log_det values should be zero tensors with the correct shape
+    assert tuple(log_det.shape) == input_shape[:-1]
+    assert tuple(log_det_inv.shape) == input_shape[:-1]
+
+
+def test_orthogonal_permutation_build_and_call(input_tensor):
+    layer = OrthogonalPermutation()
+    input_shape = keras.ops.shape(input_tensor)
+    layer.build(input_shape)
+
+    z, log_det = layer(input_tensor)
+    x_inv, log_det_inv = layer(z, inverse=True)
+
+    # Check output shapes
+    assert z.shape == input_tensor.shape
+    assert x_inv.shape == input_tensor.shape
+
+    # Forward + inverse should approximately recover input (allow some numeric tolerance)
+    np.testing.assert_allclose(
+        keras.ops.convert_to_numpy(x_inv), keras.ops.convert_to_numpy(input_tensor), rtol=1e-5, atol=1e-5
+    )
+
+    # log_det should be scalar or batched scalar
+    if len(log_det.shape) > 0:
+        assert log_det.shape[0] == input_tensor.shape[0]  # batch dim
+    else:
+        assert log_det.shape == ()
+
+    # log_det_inv should be negative of log_det (det(inv) = 1/det)
+    log_det_np = keras.ops.convert_to_numpy(log_det)
+    log_det_inv_np = keras.ops.convert_to_numpy(log_det_inv)
+    np.testing.assert_allclose(log_det_inv_np, -log_det_np, rtol=1e-5, atol=1e-5)
+
+
+def test_random_permutation_build_and_call(input_tensor):
+    layer = RandomPermutation()
+    input_shape = keras.ops.shape(input_tensor)
+    layer.build(input_shape)
+
+    # Assert forward_indices and inverse_indices are set and consistent
+    fwd = keras.ops.convert_to_numpy(layer.forward_indices)
+    inv = keras.ops.convert_to_numpy(layer.inverse_indices)
+    # Applying inv on fwd must yield ordered indices
+    reordered = fwd[inv]
+    np.testing.assert_array_equal(np.arange(len(fwd)), reordered)
+
+    z, log_det = layer(input_tensor)
+    x_inv, log_det_inv = layer(z, inverse=True)
+
+    assert z.shape == input_tensor.shape
+    assert x_inv.shape == input_tensor.shape
+    np.testing.assert_allclose(keras.ops.convert_to_numpy(x_inv), keras.ops.convert_to_numpy(input_tensor), atol=1e-5)
+    assert tuple(log_det.shape) == input_shape[:-1]
+    assert tuple(log_det_inv.shape) == input_shape[:-1]
+
+
+def test_swap_build_and_call(input_tensor):
+    layer = Swap()
+    input_shape = keras.ops.shape(input_tensor)
+    layer.build(input_shape)
+
+    fwd = keras.ops.convert_to_numpy(layer.forward_indices)
+    inv = keras.ops.convert_to_numpy(layer.inverse_indices)
+    reordered = fwd[inv]
+    np.testing.assert_array_equal(np.arange(len(fwd)), reordered)
+
+    z, log_det = layer(input_tensor)
+    x_inv, log_det_inv = layer(z, inverse=True)
+
+    assert z.shape == input_tensor.shape
+    assert x_inv.shape == input_tensor.shape
+    np.testing.assert_allclose(keras.ops.convert_to_numpy(x_inv), keras.ops.convert_to_numpy(input_tensor), atol=1e-5)
+    assert tuple(log_det.shape) == input_shape[:-1]
+    assert tuple(log_det_inv.shape) == input_shape[:-1]
diff --git a/tests/test_networks/test_embeddings.py b/tests/test_networks/test_embeddings.py
diff --git a/tests/test_networks/test_summary_networks.py b/tests/test_networks/test_summary_networks.py
diff --git a/tests/utils/jupyter.py b/tests/utils/jupyter.py

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-About us`
	`1`	`+About Us`
`2`	`2`	`========`
`3`	`3`
`4`	`4`	`Core maintainers`