[no ci] Merge remote-tracking branch 'upstream/main' into docs-migration-advice

Author: vpratz

bayesflow/networks/flow_matching/flow_matching.py

@@ -39,11 +39,11 @@ class FlowMatching(InferenceNetwork):
     }
 
     OPTIMAL_TRANSPORT_DEFAULT_CONFIG = {
-        "method": "sinkhorn",
-        "cost": "euclidean",
+        "method": "log_sinkhorn",
         "regularization": 0.1,
         "max_steps": 100,
-        "tolerance": 1e-4,
+        "atol": 1e-5,
+        "rtol": 1e-4,
     }
 
     INTEGRATE_DEFAULT_CONFIG = {

bayesflow/utils/_docs/_populate_all.py

@@ -1,23 +1,31 @@
-import inspect
+import sys
+import types
 
 
 def _add_imports_to_all(include_modules: bool | list[str] = False, exclude: list[str] | None = None):
     """Add all global variables to __all__"""
     if not isinstance(include_modules, (bool, list)):
         raise ValueError("include_modules must be a boolean or a list of strings")
 
-    exclude = exclude or []
-    calling_module = inspect.stack()[1]
-    local_stack = calling_module[0]
-    global_vars = local_stack.f_globals
-    all_vars = global_vars["__all__"] if "__all__" in global_vars else []
-    included_vars = []
-    for var_name in set(global_vars.keys()):
-        if inspect.ismodule(global_vars[var_name]):
-            if include_modules is True and var_name not in exclude and not var_name.startswith("_"):
-                included_vars.append(var_name)
-            elif isinstance(include_modules, list) and var_name in include_modules:
-                included_vars.append(var_name)
-        elif var_name not in exclude and not var_name.startswith("_"):
-            included_vars.append(var_name)
-    global_vars["__all__"] = sorted(list(set(all_vars).union(included_vars)))
+    exclude_set = set(exclude or [])
+    contains = exclude_set.__contains__
+    mod_type = types.ModuleType
+    frame = sys._getframe(1)
+    g: dict = frame.f_globals
+    existing = set(g.get("__all__", []))
+
+    to_add = []
+    include_list = include_modules if isinstance(include_modules, list) else ()
+    inc_all = include_modules is True
+
+    for name, val in g.items():
+        if name.startswith("_") or contains(name):
+            continue
+
+        if isinstance(val, mod_type):
+            if inc_all or name in include_list:
+                to_add.append(name)
+        else:
+            to_add.append(name)
+
+    g["__all__"] = sorted(existing.union(to_add))

bayesflow/utils/optimal_transport/log_sinkhorn.py

@@ -1,7 +1,6 @@
 import keras
 
 from .. import logging
-from ..tensor_utils import is_symbolic_tensor
 
 from .euclidean import euclidean
 
@@ -27,9 +26,6 @@ def log_sinkhorn_plan(x1, x2, regularization: float = 1.0, rtol=1e-5, atol=1e-8,
 
     log_plan = cost / -(regularization * keras.ops.mean(cost) + 1e-16)
 
-    if is_symbolic_tensor(log_plan):
-        return log_plan
-
     def contains_nans(plan):
         return keras.ops.any(keras.ops.isnan(plan))
 
@@ -59,7 +55,7 @@ def do_nothing():
     def log_steps():
         msg = "Log-Sinkhorn-Knopp converged after {:d} steps."
 
-        logging.info(msg, steps)
+        logging.debug(msg, steps)
 
     def warn_convergence():
         marginals = keras.ops.logsumexp(log_plan, axis=0)

bayesflow/utils/serialization.py

@@ -4,6 +4,7 @@
 import inspect
 import keras
 import numpy as np
+import sys
 
 # this import needs to be exactly like this to work with monkey patching
 from keras.saving import deserialize_keras_object
@@ -97,7 +98,10 @@ def deserialize(obj, custom_objects=None, safe_mode=True, **kwargs):
             # we marked this as a type during serialization
             obj = obj[len(_type_prefix) :]
             tp = keras.saving.get_registered_object(
-                obj, custom_objects=custom_objects, module_objects=builtins.__dict__ | np.__dict__
+                # TODO: can we pass module objects without overwriting numpy's dict with builtins?
+                obj,
+                custom_objects=custom_objects,
+                module_objects=np.__dict__ | builtins.__dict__,
             )
             if tp is None:
                 raise ValueError(
@@ -117,10 +121,9 @@ def deserialize(obj, custom_objects=None, safe_mode=True, **kwargs):
 @allow_args
 def serializable(cls, package=None, name=None):
     if package is None:
-        # get the calling module's name, e.g. "bayesflow.networks.inference_network"
-        stack = inspect.stack()
-        module = inspect.getmodule(stack[1][0])
-        package = copy(module.__name__)
+        frame = sys._getframe(1)
+        g = frame.f_globals
+        package = g.get("__name__", "bayesflow")
 
     if name is None:
         name = copy(cls.__name__)

docsrc/source/index.md

@@ -67,13 +67,15 @@ More tutorials are always welcome! Please consider making a pull request if you
 
     .. tab-item:: pip
 
-        The v2 version is not available on PyPI yet, please install from source.
+        .. code-block:: bash
+
+            pip install bayesflow
 
     .. tab-item:: source
 
         .. code-block:: bash
 
-            pip install git+https://github.com/bayesflow-org/bayesflow.git
+            pip install git+https://github.com/bayesflow-org/bayesflow.git@dev
 ```
 
 ### Backend

examples/Linear_Regression_Starter.ipynb

@@ -28,30 +28,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2025-02-14T10:51:27.573003Z",
      "start_time": "2025-02-14T10:51:27.568939Z"
     }
    },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "WARNING:bayesflow:\n",
-      "When using torch backend, we need to disable autograd by default to avoid excessive memory usage. Use\n",
-      "\n",
-      "with torch.enable_grad():\n",
-      "    ...\n",
-      "\n",
-      "in contexts where you need gradients (e.g. custom training loops).\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "import numpy as np\n",
+    "from pathlib import Path\n",
     "\n",
     "import keras\n",
     "import bayesflow as bf"
@@ -598,7 +585,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": null,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2025-02-14T10:52:51.132695Z",
@@ -618,7 +605,7 @@
     }
    ],
    "source": [
-    "f = bf.diagnostics.plots.loss(history, )"
+    "f = bf.diagnostics.plots.loss(history)"
    ]
   },
   {
@@ -964,30 +951,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 30,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "2025-04-21 11:54:04.969579: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
-      "2025-04-21 11:54:04.977366: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
-      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
-      "E0000 00:00:1745250844.984817 4140753 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
-      "E0000 00:00:1745250844.987174 4140753 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
-      "W0000 00:00:1745250844.993850 4140753 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
-      "W0000 00:00:1745250844.993860 4140753 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
-      "W0000 00:00:1745250844.993861 4140753 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
-      "W0000 00:00:1745250844.993863 4140753 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.\n",
-      "2025-04-21 11:54:04.996047: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
-      "To enable the following instructions: AVX2 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "# Recommended - full serialization (checkpoints folder must exist)\n",
-    "workflow.approximator.save(filepath=\"checkpoints/regression.keras\")\n",
+    "filepath = Path(\"checkpoints\") / \"regression.keras\"\n",
+    "filepath.parent.mkdir(exist_ok=True)\n",
+    "workflow.approximator.save(filepath=filepath)\n",
     "\n",
     "# Not recommended due to adapter mismatches - weights only\n",
     "# approximator.save_weights(filepath=\"checkpoints/regression.h5\")"
@@ -1002,21 +973,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 31,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/radevs/anaconda3/envs/bf/lib/python3.11/site-packages/keras/src/saving/serialization_lib.py:734: UserWarning: `compile()` was not called as part of model loading because the model's `compile()` method is custom. All subclassed Models that have `compile()` overridden should also override `get_compile_config()` and `compile_from_config(config)`. Alternatively, you can call `compile()` manually after loading.\n",
-      "  instance.compile_from_config(compile_config)\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "# Load approximator\n",
-    "approximator = keras.saving.load_model(\"checkpoints/regression.keras\")"
+    "approximator = keras.saving.load_model(filepath)"
    ]
   },
   {
@@ -1052,13 +1014,6 @@
     "    variable_names=par_names\n",
     ")"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
@@ -1073,16 +1028,7 @@
    "name": "python3"
   },
   "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.11"
+   "name": "python"
   },
   "widgets": {
    "application/vnd.jupyter.widget-state+json": {

examples/Lotka_Volterra_Point_Estimation_and_Expert_Stats.ipynb

@@ -37,6 +37,7 @@
    "source": [
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
+    "from pathlib import Path\n",
     "import seaborn as sns\n",
     "\n",
     "import scipy\n",
@@ -748,7 +749,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "checkpoint_path = \"checkpoints/model.keras\"\n",
+    "checkpoint_path = Path(\"checkpoints\") / \"model.keras\"\n",
+    "checkpoint_path.parent.mkdir(exist_ok=True)\n",
     "keras.saving.save_model(point_inference_workflow.approximator, checkpoint_path)"
    ]
   },

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "bayesflow"
-version = "2.0.0"
+version = "2.0.1"
 authors = [{ name = "The BayesFlow Team" }]
 classifiers = [
     "Development Status :: 5 - Production/Stable",
@@ -37,6 +37,8 @@ all = [
     "jupyter",
     "jupyterlab",
     "nbconvert",
+    "ipython",
+    "ipykernel",
     "pre-commit",
     "ruff",
     "tox",
@@ -72,6 +74,8 @@ docs = [
 ]
 test = [
     "nbconvert",
+    "ipython",
+    "ipykernel",
     "pytest",
     "pytest-cov",
     "pytest-rerunfailures",

tests/test_examples/test_examples.py

@@ -3,6 +3,7 @@
 from tests.utils import run_notebook
 
 
+@pytest.mark.skip(reason="requires setting up Stan")
 @pytest.mark.slow
 def test_bayesian_experimental_design(examples_path):
     run_notebook(examples_path / "Bayesian_Experimental_Design.ipynb")
@@ -30,7 +31,7 @@ def test_one_sample_ttest(examples_path):
 
 @pytest.mark.slow
 def test_sir_posterior_estimation(examples_path):
-    run_notebook(examples_path / "SIR_Posterior_estimation.ipynb")
+    run_notebook(examples_path / "SIR_Posterior_Estimation.ipynb")
 
 
 @pytest.mark.slow

tests/test_networks/conftest.py

@@ -85,6 +85,7 @@ def typical_point_inference_network_subnet():
         "spline_coupling_flow",
         "flow_matching",
         "free_form_flow",
+        "consistency_model",
     ],
     scope="function",
 )
@@ -106,7 +107,8 @@ def inference_network_subnet(request):
 
 
 @pytest.fixture(
-    params=["affine_coupling_flow", "spline_coupling_flow", "flow_matching", "free_form_flow"], scope="function"
+    params=["affine_coupling_flow", "spline_coupling_flow", "flow_matching", "free_form_flow", "consistency_model"],
+    scope="function",
 )
 def generative_inference_network(request):
     return request.getfixturevalue(request.param)