BayesFlow

An amortized Bayesian workflow replaces slow, per-dataset inference (like MCMC) with a fast, reusable neural inference engine. Instead of solving a new inference problem every time new data arrives, you invest upfront in a three-step process:

1. Generative Model: A simulator that produces parameters, latent variables, and observations.
2. Data Representation Pipeline: A mechanism to convert raw simulator output into structured tensors for neural networks.
3. Neural Inference Engine: A trained network that learns the inverse mapping from observations back to Bayesian quantities (e.g., posterior samples, likelihoods, or point estimates).

Once trained, the inference becomes nearly instantaneous and can be reused for any new data generated by the same model family.

BayesFlow uses a stage parameter to track the phase of functionality execution. This parameter ensures that stateful objects, such as approximators and adapters, behave correctly according to whether you are training, validating, or performing inference. You should specify the appropriate stage to prevent unintended state changes during validation or inference.

# Example conceptual usage of the stage parameter
# (Note: exact API syntax depends on the specific method being called)

# During training: allows updating the approximator
result = model.call(..., stage='training')

# During validation: identical to training but prevents updates to the approximator
result = model.call(..., stage='validation')

# During inference: prevents updates and may use a different input structure
result = model.call(..., stage='inference')

BayesFlow defines three distinct stages for its operations:

1. training: Used to train the approximator and related stateful objects (such as the adapter). This stage allows for parameter updates.
2. validation: Uses the same settings as the training stage, but calls made in this stage must not change the approximator.
3. inference: Calls in this stage must not change the approximator. The input structure may differ from the training phase. For example, during sampling, you provide summary_conditions and inference_conditions, but you do not provide the inference_variables (as these are the variables being inferred).

Starting from version 2, BayesFlow is built on Keras 3. This architecture allows you to write machine learning pipelines that are backend-agnostic, meaning the same code can run using JAX, TensorFlow, or PyTorch.

BayesFlow leverages Keras functionality and extends it with backend-specific code only when necessary to maintain this portability across different machine learning frameworks.

BayesFlow uses a serialization strategy based on storing constructor arguments (__init__ arguments) to reconstruct objects. This allows for saving the structure of a class and then mapping stored weights/state to the newly constructed object.

To support serialization, all constructor arguments must be either basic Python objects (e.g., int, float, str, bool) or other serializable objects. If you use non-serializable types, you must manually implement serialization logic.

BayesFlow relies on Keras, which can use different backends (JAX, PyTorch, or TensorFlow). By default, BayesFlow attempts to automatically select a backend based on what is installed in your environment, prioritizing them in this order: JAX, PyTorch, and then TensorFlow.

To manually select a specific backend, you must set the KERAS_BACKEND environment variable before importing bayesflow in your Python script.

Important Note for PyTorch users: When using the torch backend, BayesFlow disables torch.autograd by default to prevent excessive memory usage. If you are implementing custom training loops or any logic that requires gradients, you must explicitly wrap that code in a with torch.enable_grad(): block.

import os
# Set the backend before importing bayesflow
os.environ["KERAS_BACKEND"] = "jax"

import bayesflow

BayesFlow uses Keras3 for multi-backend support, allowing you to use JAX, PyTorch, or TensorFlow. JAX is recommended for performance.

As of version 2.0.7, the backend is set automatically if only one is installed. If you have multiple backends, you must set the KERAS_BACKEND environment variable.

Important: The environment variable must be set before importing bayesflow.

import os
# Set the backend before importing bayesflow
os.environ["KERAS_BACKEND"] = "jax"
import bayesflow