MeanFlow

class bayesflow.experimental.MeanFlow(*args, **kwargs)

Bases: InferenceNetwork

Mean flow for few-step sampling and density evaluation in simulation-based inference.

Learns the average velocity u(z, s, t) of the flow matching path, which defines the flow map z_t = z_s + (t - s) * u(z_s, s, t) between any two times [1, 2], trained to be invertible [3]. Time t = 0 is the data and t = 1 the base distribution. Sampling and density evaluation compose steps flow maps on an EDM time schedule, and the density uses the log-determinant of each map, or with density_head=True a learned average divergence of each map [4].

Note: With the PyTorch backend without CUDA, subnet="diffusion_transformer" is not supported because PyTorch scaled-dot-product attention on CPU and MPS lacks forward-mode AD support.

Parameters:

subnetstr, type, or keras.Layer

A neural network type for the mean flow model, will be instantiated using subnet_kwargs. If a string is provided, it should be a registered name (e.g., "time_mlp"). If a type or keras.Layer is provided, it will be directly instantiated with the given subnet_kwargs. Any subnet must accept a tuple of tensors (target, time, conditions), where time holds the two times (s, s - t).

base_distributionstr or Distribution

The base probability distribution from which samples are drawn. Default is "normal".

use_optimal_transportbool

Whether to apply optimal transport for improved training stability. Default is False.

loss_fnstr or keras.Loss

The loss function of the average velocity, summed over the target dimensions, and of the invertibility loss. Default is "mse".

integrate_kwargsdict[str, Any], optional

Additional keyword arguments for the ODE integrator used with use_instantaneous_flow=True.

optimal_transport_kwargsdict[str, Any], optional

Additional keyword arguments for configuring optimal transport.

subnet_kwargsdict[str, Any], optional

Keyword arguments passed to the subnet constructor or used to update the default MLP settings.

time_power_law_alphafloat

Changes the distribution of sampled times during training. Time is sampled as t = u^(1+alpha) with u ~ U(0, 1), i.e., p(t) ~ t^(-alpha/(1+alpha)). Positive values emphasize times near the data endpoint at t = 0; alpha = 0 is uniform. Default is 0.5.

stepsint

Default number of flow maps used for sampling and density evaluation. Can be overridden per call by passing steps= to sampling. Default is 4.

velocity_head_layersint

Number of final subnet blocks that are duplicated into a head predicting the instantaneous velocity. Requires a subnet with embed, compute_embed_output_shape and blocks, such as "time_mlp" or "diffusion_transformer". With 0, the instantaneous velocity is the average velocity at s = t. Default is 2.

velocity_loss_fnstr or keras.Loss

The loss function of the velocity head. Default is "mse".

density_headbool

Whether to learn the average divergence of each flow map as a second output of the subnet, which then replaces the exact log-determinants in density evaluation. Default is False.

boundary_probfloat

Fraction of training time pairs with s = t. Default is 0.7.

cycle_weightfloat

Weight of the invertibility loss, which maps from s to t and back. Default is 2.5.

cycle_mismatch_probfloat

Probability of evaluating the invertibility loss at the conditions of another sample in the batch. Default is 0.6.

adaptive_loss_powerfloat

Power p of the adaptive weight (loss + adaptive_loss_eps)^(-p) of the average velocity loss [2], 0 disables the weighting. Default is 0.25.

adaptive_loss_epsfloat

Constant of the adaptive loss weight. Default is 0.1.

mask_handler~bayesflow.masking.MaskHandler, optional

Handler used to generate and route masks. If None, a default mask handler is created.

**kwargs

Additional keyword arguments passed to the base InferenceNetwork.

References

[1] Geng et al. (2025). Mean flows for one-step generative modeling. NeurIPS, 38. [2] Geng et al. (2026). Improved mean flows: On the challenges of fastforward generative models. CVPR. [3] Rehman et al. (2026). FALCON: Few-step accurate likelihoods for continuous flows. ICLR. [4] OuYang et al. (2026). Few-step Boltzmann generators via scalable likelihood flow maps. arXiv:2606.29110.

build(xz_shape: tuple[int, ...], conditions_shape: tuple[int, ...] = None) → None

get_config()

Returns the config of the object.

An object config is a Python dictionary (serializable) containing the information needed to re-instantiate it.

compute_metrics(x: Tensor, conditions: Tensor = None, sample_weight: Tensor = None, stage: str = 'training', **kwargs) → dict[str, Tensor]

velocity(xz: Tensor, time: float | Tensor, conditions: Tensor = None, training: bool = False, **kwargs) → Tensor

static make_time_schedule(steps: int, rho: float = 7.0, sigma_min: float = 0.001) → Tensor

EDM time schedule of steps steps, decreasing from 1 to 0.

__call__(*args, **kwargs)

Call self as a function.

add_loss(loss)

Can be called inside of the call() method to add a scalar loss.

Example:

class MyLayer(Layer):
    ...
    def call(self, x):
        self.add_loss(ops.sum(x))
        return x

add_metric(*args, **kwargs)

add_variable(shape, initializer, dtype=None, trainable=True, autocast=True, regularizer=None, constraint=None, name=None)

Add a weight variable to the layer.

Alias of add_weight().

add_weight(*args, shape=None, initializer=None, dtype=None, trainable=True, autocast=True, regularizer=None, constraint=None, aggregation='none', overwrite_with_gradient=False, name=None)

Add a weight variable to the layer.

Args:

shape: Shape tuple for the variable. Must be fully-defined

(no None entries). Defaults to () (scalar) if unspecified.

initializer: Initializer object to use to populate the initial

variable value, or string name of a built-in initializer (e.g. “random_normal”). If unspecified, defaults to “glorot_uniform” for floating-point variables and to “zeros” for all other types (e.g. int, bool).

dtype: Dtype of the variable to create, e.g. “float32”. If

unspecified, defaults to the layer’s variable dtype (which itself defaults to “float32” if unspecified).

trainable: Boolean, whether the variable should be trainable via

backprop or whether its updates are managed manually. Defaults to True.

autocast: Boolean, whether to autocast layers variables when

accessing them. Defaults to True.

regularizer: Regularizer object to call to apply penalty on the

weight. These penalties are summed into the loss function during optimization. Defaults to None.

constraint: Contrainst object to call on the variable after any

optimizer update, or string name of a built-in constraint. Defaults to None.

aggregation: Optional string, one of None, “none”, “mean”,

“sum” or “only_first_replica”. Annotates the variable with the type of multi-replica aggregation to be used for this variable when writing custom data parallel training loops. Defaults to “none”.

overwrite_with_gradient: Boolean, whether to overwrite the variable

with the computed gradient. This is useful for float8 training. Defaults to False.

name: String name of the variable. Useful for debugging purposes.

build_from_config(config)

Builds the layer’s states with the supplied config dict.

By default, this method calls the build(config[“input_shape”]) method, which creates weights based on the layer’s input shape in the supplied config. If your config contains other information needed to load the layer’s state, you should override this method.

Args:

config: Dict containing the input shape associated with this layer.

call(xz: Tensor, conditions: Tensor = None, inverse: bool = False, density: bool = False, training: bool = False, seed: int | SeedGenerator | None = None, hutchinson_steps: int | None = None, guidance_kwargs: Mapping[str, Any] | None = None, **kwargs) → Tensor | tuple[Tensor, Tensor]

property compute_dtype

The dtype of the computations performed by the layer.

compute_mask(inputs, previous_mask)

compute_output_shape(*args, **kwargs)

compute_output_spec(*args, **kwargs)

count_params()

Count the total number of scalars composing the weights.

Returns:

An integer count.

property dtype

Alias of layer.variable_dtype.

property dtype_policy

classmethod from_config(config, custom_objects=None)

Creates an operation from its config.

This method is the reverse of get_config, capable of instantiating the same operation from the config dictionary.

Note: If you override this method, you might receive a serialized dtype config, which is a dict. You can deserialize it as follows:

if "dtype" in config and isinstance(config["dtype"], dict):
    policy = dtype_policies.deserialize(config["dtype"])

Args:

config: A Python dictionary, typically the output of get_config.

Returns:

An operation instance.

get_build_config()

Returns a dictionary with the layer’s input shape.

This method returns a config dict that can be used by build_from_config(config) to create all states (e.g. Variables and Lookup tables) needed by the layer.

By default, the config only contains the input shape that the layer was built with. If you’re writing a custom layer that creates state in an unusual way, you should override this method to make sure this state is already created when Keras attempts to load its value upon model loading.

Returns:

A dict containing the input shape associated with the layer.

get_weights()

Return the values of layer.weights as a list of NumPy arrays.

property input

Retrieves the input tensor(s) of a symbolic operation.

Only returns the tensor(s) corresponding to the first time the operation was called.

Returns:

Input tensor or list of input tensors.

property input_dtype

The dtype layer inputs should be converted to.

property input_spec

load_own_variables(store)

Loads the state of the layer.

You can override this method to take full control of how the state of the layer is loaded upon calling keras.models.load_model().

Args:

store: Dict from which the state of the model will be loaded.

log_prob(samples: Tensor, conditions: Tensor = None, **kwargs) → Tensor

property losses

List of scalar losses from add_loss, regularizers and sublayers.

property metrics

List of all metrics.

property metrics_variables

List of all metric variables.

property non_trainable_variables

List of all non-trainable layer state.

This extends layer.non_trainable_weights to include all state used by the layer including state for metrics and `SeedGenerator`s.

property non_trainable_weights

List of all non-trainable weight variables of the layer.

These are the weights that should not be updated by the optimizer during training. Unlike, layer.non_trainable_variables this excludes metric state and random seeds.

property output

Retrieves the output tensor(s) of a layer.

Only returns the tensor(s) corresponding to the first time the operation was called.

Returns:

Output tensor or list of output tensors.

property path

The path of the layer.

If the layer has not been built yet, it will be None.

property quantization_mode

The quantization mode of this layer, None if not quantized.

quantize(mode=None, type_check=True, config=None)

quantized_build(input_shape, mode)

quantized_call(*args, **kwargs)

rematerialized_call(layer_call, *args, **kwargs)

Enable rematerialization dynamically for layer’s call method.

Args:

layer_call: The original call method of a layer.

Returns:

Rematerialized layer’s call method.

sample(batch_shape: tuple[int, ...], conditions: Tensor = None, seed: int | SeedGenerator | None = None, **kwargs) → Tensor

save_own_variables(store)

Saves the state of the layer.

You can override this method to take full control of how the state of the layer is saved upon calling model.save().

Args:

store: Dict where the state of the model will be saved.

set_weights(weights)

Sets the values of layer.weights from a list of NumPy arrays.

stateless_call(trainable_variables, non_trainable_variables, *args, return_losses=False, **kwargs)

Call the layer without any side effects.

Args:

trainable_variables: List of trainable variables of the model. non_trainable_variables: List of non-trainable variables of the

model.

*args: Positional arguments to be passed to call(). return_losses: If True, stateless_call() will return the list of

losses created during call() as part of its return values.

**kwargs: Keyword arguments to be passed to call().

Returns:

A tuple. By default, returns (outputs, non_trainable_variables).

If return_losses = True, then returns (outputs, non_trainable_variables, losses).

Note: non_trainable_variables include not only non-trainable weights such as BatchNormalization statistics, but also RNG seed state (if there are any random operations part of the layer, such as dropout), and Metric state (if there are any metrics attached to the layer). These are all elements of state of the layer.

Example:

model = ...
data = ...
trainable_variables = model.trainable_variables
non_trainable_variables = model.non_trainable_variables
# Call the model with zero side effects
outputs, non_trainable_variables = model.stateless_call(
    trainable_variables,
    non_trainable_variables,
    data,
)
# Attach the updated state to the model
# (until you do this, the model is still in its pre-call state).
for ref_var, value in zip(
    model.non_trainable_variables, non_trainable_variables
):
    ref_var.assign(value)

property supports_masking

Whether this layer supports computing a mask using compute_mask.

symbolic_call(*args, **kwargs)

property trainable

Settable boolean, whether this layer should be trainable or not.

property trainable_variables

List of all trainable layer state.

This is equivalent to layer.trainable_weights.

property trainable_weights

List of all trainable weight variables of the layer.

These are the weights that get updated by the optimizer during training.

property variable_dtype

The dtype of the state (weights) of the layer.

property variables

List of all layer state, including random seeds.

This extends layer.weights to include all state used by the layer including `SeedGenerator`s.

Note that metrics variables are not included here, use metrics_variables to visit all the metric variables.

property weights

List of all weight variables of the layer.

Unlike, layer.variables this excludes metric state and random seeds.