MaskHandler#

class bayesflow.masking.MaskHandler(*, fixed_target_prob: float = 0.0, missing_target_prob: float = 0.0, missing_conditions_prob: float = 0.0, seed: int | None = None)[source]#

Bases: object

Generate semantic masks and route masks between BayesFlow interfaces.

Semantic masks change the target distribution. For example, masking missing targets lets a network trained on \(p(x_1, x_2)\) learn its marginals. Padding and attention masks are structural; they support operations such as learning filtering distributions with an AutoregressiveEstimator.

Parameters:

fixed_target_probfloat, optional: Probability of fixing a target coordinate during training.
missing_target_probfloat, optional: Probability of hiding a fixed target coordinate during training.
missing_conditions_probfloat, optional: Probability of hiding a condition coordinate during training.
seedint, optional: Seed used to initialize the local random generator for sampling training masks.

summary_mask_routes = {MaskName.SUMMARY_ATTENTION: MaskName.ATTENTION_MASK, MaskName.SUMMARY: MaskName.MASK}#

inference_mask_routes = {MaskName.INFERENCE_ATTENTION: MaskName.ATTENTION_MASK, MaskName.INFERENCE: MaskName.MASK}#

subnet_mask_names = (MaskName.ATTENTION_MASK, MaskName.MASK, MaskName.FIXED_TARGET, MaskName.INFER_TARGET, MaskName.OBSERVED_CONDITION)#

get_config() → dict[str, Any][source]#

classmethod from_config(config: dict[str, Any]) → MaskHandler[source]#

classmethod collect_summary_masks(*sources: Mapping[str, Any]) → dict[str, Any][source]#: Map summary-boundary mask names to their subnet names.

classmethod collect_inference_masks(*sources: Mapping[str, Any]) → dict[str, Any][source]#: Map inference-boundary mask names to their subnet names.

classmethod collect_semantic_masks(*sources: Mapping[str, Any]) → dict[str, Any][source]#: Collect semantic masks without renaming them.

classmethod collect_mask_kwargs(*sources: Mapping[str, Any]) → dict[str, Any][source]#: Collect canonical subnet masks and the fixed target value.

classmethod collect_subnet_masks(subnet: Layer, *sources: Mapping[str, Any]) → dict[str, Any][source]#: Collect masks explicitly named in subnet.call.

classmethod with_inference_masks(kwargs: Mapping[str, Any], adapted: Mapping[str, Any] | None = None) → dict[str, Any][source]#

Add masks discovered in adapted data to inference call kwargs.

Explicit kwargs take precedence over masks found in adapted data. Boundary names such as inference_mask are replaced by the canonical subnet names mask and attention_mask.

classmethod without_semantic_kwargs(source: Mapping[str, Any]) → dict[str, Any][source]#: Return kwargs without semantic masks or fixed target values.

classmethod without_mask_kwargs(source: Mapping[str, Any]) → dict[str, Any][source]#: Return kwargs without canonical subnet or semantic mask fields.

classmethod reject_structural_masks(*sources: Mapping[str, Any], reason: str) → None[source]#: Raise if an attention or padding mask would be ignored.

static maybe_apply(value: Tensor, mask: Tensor | None, replacement: Tensor | None = None) → Tensor[source]#: Keep value where mask is one and use replacement where it is zero. If no replacement is provided, a zero replacement will be used.

static per_token_target_mask(mask: Tensor | None, targets: Tensor) → Tensor | None[source]#: Reshape a (num_targets,) or (batch, num_targets) target mask to (batch, num_targets, 1), one value per target token of a transformer that broadcasts over the token width.

classmethod has_batch_axis(name: str, value: Tensor, batch_size: int | None, *, target_structural_rank: int = 0, condition_structural_rank: int = 0) → bool[source]#

Return whether value carries the leading data-batch axis.

Padding masks use shared (S,) and batched (B, S) shapes; attention masks use shared (T, S) and batched (B, T, S) shapes. Semantic masks additionally account for the structure of the tensor they mask. If a shape is also valid without a batch axis, a matching leading batch dimension takes precedence.

classmethod batched_keys(source: Mapping[str, Any], batch_size: int | None, target_structural_rank: int = 0, condition_structural_rank: int = 0) → set[str][source]#: Return tensor fields that carry the leading data-batch axis.

classmethod infer_batch_size(source: Mapping[str, Any], target_structural_rank: int = 0) → int | None[source]#: Infer the batch size from a semantic target mask or fixed value.

classmethod select_autoregressive_step(name: str, value: Tensor, *, step: int, num_steps: int, has_batch_axis: bool) → Tensor[source]#

Select one autoregressive step when value has a step axis.

Shared step axes use axis 0 and batched step axes use axis 1. Singleton step axes are reused, while event-only values are returned unchanged.

resolve(*sources: Mapping[str, Any], targets: Tensor, conditions: Tensor | None = None, stage: str = 'inference', subnet: Layer | None = None) → InferenceMasks[source]#

Resolve semantic masks for one inference-network operation.

Parameters:

*sourcesMapping[str, Any]: Mappings containing explicit semantic masks. Later mappings take precedence over earlier ones.
targetsTensor: Target tensor whose shape is used to broadcast target masks and fixed target values.
conditionsTensor, optional: Condition tensor whose shape is used to broadcast the observed condition mask.
stagestr, optional: Current execution stage. Probability-based masks are sampled only during "training"; other stages use explicit masks only.
subnetkeras.Layer, optional: Mask consumer. Generated infer_target_mask and observed_condition_mask values are limited to masks supported by the subnet. subnet_kwargs in the result contains the masks declared by its call signature.

Returns:

InferenceMasks: Broadcast semantic masks and the masks to pass to it, if subnet is supplied.

Raises:

ValueError: If subnet is supplied and an attention_mask or mask in sources is not named in its call signature.

resolve_compositional(subnet: Layer, *sources: Mapping[str, Any], targets: Tensor, conditions: Tensor) → InferenceMasks[source]#: Resolve compositional masks, defaulting every condition to observed.

repeat_compositional(masks: InferenceMasks, num_repeats: int) → InferenceMasks[source]#

Repeat resolved masks before flattening compositional inputs.

Semantic masks must already be broadcast to their target or condition tensors. The condition mask is attached after compositional selection.

static with_flattened_condition_mask(masks: InferenceMasks, observed_condition_mask: Tensor) → InferenceMasks[source]#: Flatten and attach a selected compositional condition mask.

integrator_callbacks(masks: InferenceMasks, replacement: Tensor | None = None, state_key: str = 'xz') → tuple[Callable | None, Callable | None][source]#: Build state and update projections that preserve fixed targets.

random_mask(reference: Tensor, *, drop_prob: float, keep_one: bool = False, eligible: Tensor | None = None) → Tensor[source]#

Sample a binary keep mask with the shape of reference.

eligible limits coordinates that may be kept. If keep_one is true, at least one eligible coordinate on the final axis remains one whenever that axis contains an eligible coordinate.