Merge branch 'main' into dev

Author: LarsKue

CONTRIBUTING.md

@@ -56,30 +56,49 @@ def prior3(a, b, c):
 
 ### 2. Set up your development environment
 
-Create a fork of the BayesFlow git repository using the GitHub interface.
-Then clone your fork and install the development environment with conda:
+Once we agree on an approach in the issue you opened, we can move ahead with the implementation.
+
+First, create a development environment with conda, or any other environment manager of your choice.
+
+```bash
+conda create -n bf python=3.11
+conda activate bf
+```
+
+Note: always use a clean environment dedicated for development of BayesFlow to avoid dependency issues.
+
+Most contributors will have to create and clone a fork of BayesFlow using the GitHub interface or the CLI:
+
+```bash
+gh repo fork bayesflow-org/bayesflow --clone
+```
+
+Then, check out the development branch and install dependencies:
 
 ```bash
-git clone https://github.com/<your-username>/bayesflow
 cd bayesflow
 git checkout dev
-conda env create --file environment.yaml --name bayesflow
-conda activate bayesflow
+conda install pip
+pip install -e .[dev,docs,test]
 pre-commit install
 ```
 
-We recommend using the PyTorch backend for development.
-Be careful not to downgrade your keras version when installing the backend.
+Finally, install at least one backend of your choice.
+At the moment of writing this, to install all three backends on a machine supporting CUDA 12.6, we would use:
+
+```bash
+pip install -U jax[cuda12]
+pip install -U tensorflow[and-cuda]
+pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126
+```
 
-You can also install a triple backend cuda environment using `cuda.yaml`:
+You can set an environment variable to choose the default backend. We recommend defaulting to jax:
 
 ```bash
-conda env create --file cuda.yaml --name bf-cuda
+conda env config vars set KERAS_BACKEND=jax
 ```
 
-Note that we cannot guarantee that this environment file will always work,
-as it is not tested frequently or across devices.
-It is provided merely as a convenience. For reliability, install only a single backend.
+Note that you will have to re-activate the environment for the changes to take effect.
 
 ### 3. Implement your changes
 
@@ -138,7 +157,8 @@ z = keras.ops.convert_to_numpy(x)
 
 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.
 
-Run the following command to install all necessary packages for setting up documentation generation:
+If you haven't done so earlier, run the following command to install all necessary packages for setting up
+documentation generation:
 
 ```
 pip install .[docs]

README.md

@@ -63,18 +63,25 @@ More tutorials are always welcome! Please consider making a pull request if you
 
 ## Install
 
+BayesFlow is available to install via pip:
+
+```bash
+pip install bayesflow
+```
+
 ### Backend
 
-First, install one machine learning backend of choice. Note that BayesFlow **will not run** without a backend.
+To use BayesFlow, you will also need to install one of the following machine learning backends.
+Note that BayesFlow **will not run** without a backend.
 
 - [Install JAX](https://jax.readthedocs.io/en/latest/installation.html)
 - [Install PyTorch](https://pytorch.org/get-started/locally/)
 - [Install TensorFlow](https://www.tensorflow.org/install)
 
-If you don't know which backend to use, we recommend JAX as it is currently
-the fastest backend.
+If you don't know which backend to use, we recommend JAX as it is currently the fastest backend.
 
-Once installed, [set the backend environment variable as required by keras](https://keras.io/getting_started/#configuring-your-backend). For example, inside your Python script write:
+Once installed, [set the backend environment variable as required by keras](https://keras.io/getting_started/#configuring-your-backend).
+For example, inside your Python script write:
 
 ```python
 import os
@@ -88,31 +95,19 @@ If you use conda, you can alternatively set this individually for each environme
 conda env config vars set KERAS_BACKEND=jax
 ```
 
-This way, you also don't have to manually set the backend every time you are starting Python to use BayesFlow.
-
-**Caution:** Some people report that the IDE (e.g., VSCode or PyCharm) can silently overwrite environment variables. If you have set your backend as an environment variable and you still get keras-related import errors when loading BayesFlow, these IDE shenanigans might be the culprit. Try setting the keras backend in your Python script via `import os; os.environ["KERAS_BACKEND"] = "<YOUR-BACKEND>"`.
-
-### Using pip
-
-You can install the Bayesflow from Github with pip:
+Or just plainly set the environment variable in your shell:
 
 ```bash
-pip install git+https://github.com/bayesflow-org/bayesflow@main
+export KERAS_BACKEND=jax
 ```
 
-### Using Conda
+This way, you also don't have to manually set the backend every time you are starting Python to use BayesFlow.
 
-Bayesflow is currently not conda-installable.
+**Caution:** Some development environments (e.g., VSCode or PyCharm) can silently overwrite environment variables. If you have set your backend as an environment variable and you still get keras-related import errors when loading BayesFlow, these IDE shenanigans might be the culprit. Try setting the keras backend in your Python script via `import os; os.environ["KERAS_BACKEND"] = "<YOUR-BACKEND>"`.
 
 ### From Source
 
-If you want to contribute to BayesFlow, we recommend installing it from source:
-
-```bash
-git clone -b main [email protected]:bayesflow-org/bayesflow.git
-cd <local-path-to-bayesflow-repository>
-conda env create --file environment.yaml --name bayesflow
-```
+If you want to contribute to BayesFlow, we recommend installing it from source, see [CONTRIBUTING.md](CONTRIBUTING.md) for more details.
 
 ## Reporting Issues
 
@@ -170,17 +165,54 @@ You can cite BayesFlow along the lines of:
 
 ## FAQ
 
-- *I am starting with Bayesflow, which backend shall I use?*
-**A**: We recommend JAX as it is currently the fastest backend.
+-------------
+
+**Question:**
+I am starting with Bayesflow, which backend should I use?
 
-- *What is the difference between Bayesflow 2.0+ and previous versions?*
-**A**: BayesFlow 2.0+ is a complete rewrite of the library. It shares the same
+**Answer:**
+We recommend JAX as it is currently the fastest backend.
+
+-------------
+
+**Question:**
+I am getting `ModuleNotFoundError: No module named 'tensorflow'` when I try to import BayesFlow.
+
+**Answer:**
+One of these applies:
+- You want to use tensorflow as your backend, but you have not installed it.
+See [here](https://www.tensorflow.org/install).
+
+
+- You want to use a backend other than tensorflow, but have not set the environment variable correctly.
+See [here](https://keras.io/getting_started/#configuring-your-backend).
+
+
+- You have set the environment variable, but it is not being picked up by Python.
+This can happen silently in some development environments (e.g., VSCode or PyCharm).
+Try setting the backend as shown [here](https://keras.io/getting_started/#configuring-your-backend)
+in your Python script via `os.environ`.
+
+-------------
+
+**Question:**
+What is the difference between Bayesflow 2.0+ and previous versions?
+
+**Answer:**
+BayesFlow 2.0+ is a complete rewrite of the library. It shares the same
 overall goals with previous versions, but has much better modularity
 and extensibility. What is more, the new BayesFlow has multi-backend support via Keras3,
 while the old version was based on TensorFlow.
 
-- *I still need the old BayesFlow for some of my projects. How can I install it?*
-**A**: You can find and install the old Bayesflow version via the `stable-legacy` branch on GitHub.
+-------------
+
+**Question:**
+I still need the old BayesFlow for some of my projects. How can I install it?
+
+**Answer:**
+You can find and install the old Bayesflow version via the `stable-legacy` branch on GitHub.
+
+-------------
 
 ## Awesome Amortized Inference
 

bayesflow/__init__.py

@@ -34,6 +34,17 @@ def setup():
 
     logging.debug(f"Using backend {keras.backend.backend()!r}")
 
+    if keras.backend.backend() == "torch":
+        import torch
+
+        torch.autograd.set_grad_enabled(False)
+
+        logging.warning(
+            "When using torch backend, we need to disable autograd by default to avoid excessive memory usage. Use\n"
+            "with torch.enable_grad():\n"
+            "in contexts where you need gradients (e.g. custom training loops)."
+        )
+
 
 # call and clean up namespace
 setup()

bayesflow/diagnostics/plots/pairs_samples.py

@@ -115,6 +115,9 @@ def _pairs_samples(
             **kwargs,
         )
 
+        # ensures that color doesn't overwrite palette
+        color = None
+
     else:
         # plot just the one set of distributions
         data_to_plot = pd.DataFrame(plot_data["estimates"], columns=variable_names)

bayesflow/utils/decorators.py

@@ -122,7 +122,7 @@ def sanitize_input_shape(fn: Callable):
     def callback(input_shape: Shape) -> Shape:
         if input_shape[0] is None:
             input_shape = list(input_shape)
-            input_shape[0] = 32
+            input_shape[0] = 1
             return tuple(input_shape)
         return input_shape
 

examples/Bayesian_Experimental_Design.ipynb

@@ -266,7 +266,7 @@
     "    def run_hmc_posterior(\n",
     "        self, y: Tensor, num_samples: int = 1000, num_chains: int = 4\n",
     "    ) -> Tensor:\n",
-    "        stan_model = CmdStanModel(stan_file=\"mm_gsn.stan\")\n",
+    "        stan_model = CmdStanModel(stan_file=\"misc/mm_gsn.stan\")\n",
     "        hmc_posterior_samples = []\n",
     "\n",
     "        for y_sample in y:\n",
@@ -764,16 +764,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.5"
+   "name": "python"
   }
  },
  "nbformat": 4,

examples/From_ABC_to_BayesFlow.ipynb

@@ -608,7 +608,7 @@
    },
    "outputs": [],
    "source": [
-    "inference_net = bf.networks.ConsistencyModel(total_stps=epochs*num_batches_per_epoch)"
+    "inference_net = bf.networks.ConsistencyModel(total_steps=epochs*num_batches_per_epoch)"
    ]
   },
   {
@@ -933,13 +933,6 @@
    "source": [
     "We have yet to fully explore the potential of amortized inference. For higher-dimensional problems, we can train a summary network jointly with the inference network, eliminating the need to manually design summary statistics as required in ABC. Additionally, the trained approximator can be seamlessly transferred to new datasets, enabling inference without retraining. Even in this simple example, we see that training the approximator required fewer simulations than running the ABC, which is particularly beneficial for computationally expensive simulators.\n"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
@@ -950,16 +943,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.8.8"
+   "name": "python"
   }
  },
  "nbformat": 4,