Robustify kwargs passing inference networks, add class variables

Author: stefanradev93

bayesflow/approximators/approximator.py

@@ -23,7 +23,7 @@ def build_adapter(cls, **kwargs) -> Adapter:
  raise NotImplementedError
 
  def build_from_data(self, data: Mapping[str, any]) -> None:
- self.compute_metrics(**data, stage="training")
+ self.compute_metrics(**filter_kwargs(data, self.compute_metrics), stage="training")
  self.built = True
 
  @classmethod

bayesflow/approximators/continuous_approximator.py

@@ -32,6 +32,8 @@ class ContinuousApproximator(Approximator):
  Additional arguments passed to the :py:class:`bayesflow.approximators.Approximator` class.
  """
 
+ SAMPLE_KEYS = ["summary_variables", "inference_conditions"]
+
  def __init__(
  self,
  *,
@@ -51,6 +53,7 @@ def build_adapter(
  inference_variables: Sequence[str],
  inference_conditions: Sequence[str] = None,
  summary_variables: Sequence[str] = None,
+ standardize: bool = True,
  sample_weight: str = None,
  ) -> Adapter:
  """Create an :py:class:`~bayesflow.adapters.Adapter` suited for the approximator.
@@ -63,9 +66,12 @@ def build_adapter(
  Names of the inference conditions in the data
  summary_variables : Sequence of str, optional
  Names of the summary variables in the data
+ standardize : bool, optional
+ Decide whether to standardize all variables, default is True
  sample_weight : str, optional
  Name of the sample weights
  """
+
  adapter = Adapter()
  adapter.to_array()
  adapter.convert_dtype("float64", "float32")
@@ -82,7 +88,9 @@ def build_adapter(
  adapter = adapter.rename(sample_weight, "sample_weight")
 
  adapter.keep(["inference_variables", "inference_conditions", "summary_variables", "sample_weight"])
- adapter.standardize(exclude="sample_weight")
+
+ if standardize:
+ adapter.standardize(exclude="sample_weight")
 
  return adapter
 
@@ -334,12 +342,18 @@ def sample(
  dict[str, np.ndarray]
  Dictionary containing generated samples with the same keys as `conditions`.
  """
+
+ # Apply adapter transforms to raw simulated / real quantities
  conditions = self.adapter(conditions, strict=False, stage="inference", **kwargs)
- # at inference time, inference_variables are estimated by the networks and thus ignored in conditions
- conditions.pop("inference_variables", None)
+
+ # Ensure only keys relevant for sampling are present in the conditions dictionary
+ conditions = {k: v for k, v in conditions.items() if k in ContinuousApproximator.SAMPLE_KEYS}
+
  conditions = keras.tree.map_structure(keras.ops.convert_to_tensor, conditions)
  conditions = {"inference_variables": self._sample(num_samples=num_samples, **conditions, **kwargs)}
  conditions = keras.tree.map_structure(keras.ops.convert_to_numpy, conditions)
+
+ # Back-transform quantities and samples
  conditions = self.adapter(conditions, inverse=True, strict=False, **kwargs)
 
  if split:

bayesflow/approximators/model_comparison_approximator.py

@@ -30,11 +30,13 @@ class ModelComparisonApproximator(Approximator):
  The network backbone (e.g, an MLP) that is used for model classification.
  The input of the classifier network is created by concatenating `classifier_variables`
  and (optional) output of the summary_network.
- summary_network: bg.networks.SummaryNetwork, optional
+ summary_network: bf.networks.SummaryNetwork, optional
  The summary network used for data summarization (default is None).
  The input of the summary network is `summary_variables`.
  """
 
+ SAMPLE_KEYS = ["summary_variables", "inference_conditions"]
+
  def __init__(
  self,
  *,
@@ -304,9 +306,13 @@ def predict(
  np.ndarray
  Predicted posterior model probabilities given `conditions`.
  """
+
+ # Apply adapter transforms to raw simulated / real quantities
  conditions = self.adapter(conditions, strict=False, stage="inference", **kwargs)
- # at inference time, model_indices are predicted by the networks and thus ignored in conditions
- conditions.pop("model_indices", None)
+
+ # Ensure only keys relevant for sampling are present in the conditions dictionary
+ conditions = {k: v for k, v in conditions.items() if k in ModelComparisonApproximator.SAMPLE_KEYS}
+
  conditions = keras.tree.map_structure(keras.ops.convert_to_tensor, conditions)
 
  output = self._predict(**conditions, **kwargs)

bayesflow/approximators/point_approximator.py

@@ -156,8 +156,10 @@ def log_prob(
 
  def _prepare_conditions(self, conditions: Mapping[str, np.ndarray], **kwargs) -> dict[str, Tensor]:
  """Adapts and converts the conditions to tensors."""
+
  conditions = self.adapter(conditions, strict=False, stage="inference", **kwargs)
- conditions.pop("inference_variables", None)
+ conditions = {k: v for k, v in conditions.items() if k in ContinuousApproximator.SAMPLE_KEYS}
+
  return keras.tree.map_structure(keras.ops.convert_to_tensor, conditions)
 
  def _apply_inverse_adapter_to_estimates(

bayesflow/networks/consistency_models/consistency_model.py

@@ -187,7 +187,7 @@ def build(self, xz_shape, conditions_shape=None):
  self.c_huber = 0.00054 * ops.sqrt(xz_shape[-1])
  self.c_huber2 = self.c_huber**2
 
- ## Calculate discretization schedule in advance
+ # Calculate discretization schedule in advance
  # The Jax compiler requires fixed-size arrays, so we have
  # to store all the discretized_times in one matrix in advance
  # and later only access the relevant entries.
@@ -213,34 +213,24 @@ def build(self, xz_shape, conditions_shape=None):
  disc = ops.convert_to_numpy(self._discretize_time(n))
  discretized_times[i, : len(disc)] = disc
  discretization_map[n] = i
+
  # Finally, we convert the vectors to tensors
  self.discretized_times = ops.convert_to_tensor(discretized_times, dtype="float32")
  self.discretization_map = ops.convert_to_tensor(discretization_map)
 
- def call(
- self,
- xz: Tensor,
- conditions: Tensor = None,
- inverse: bool = False,
- **kwargs,
- ):
- if inverse:
- return self._inverse(xz, conditions=conditions, **kwargs)
- return self._forward(xz, conditions=conditions, **kwargs)
-
- def _forward_train(self, x: Tensor, noise: Tensor, t: Tensor, conditions: Tensor = None, **kwargs) -> Tensor:
- """Forward function for training. Calls consistency function with
- noisy input
- """
+ def _forward_train(
+ self, x: Tensor, noise: Tensor, t: Tensor, conditions: Tensor = None, training: bool = False, **kwargs
+ ) -> Tensor:
+ """Forward function for training. Calls consistency function with noisy input"""
  inp = x + t * noise
- return self.consistency_function(inp, t, conditions=conditions, **kwargs)
+ return self.consistency_function(inp, t, conditions=conditions, training=training)
 
  def _forward(self, x: Tensor, conditions: Tensor = None, **kwargs) -> Tensor:
  # Consistency Models only learn the direction from noise distribution
  # to target distribution, so we cannot implement this function.
  raise NotImplementedError("Consistency Models are not invertible")
 
- def _inverse(self, z: Tensor, conditions: Tensor = None, **kwargs) -> Tensor:
+ def _inverse(self, z: Tensor, conditions: Tensor = None, training: bool = False, **kwargs) -> Tensor:
  """Generate random draws from the approximate target distribution
  using the multistep sampling algorithm from [1], Algorithm 1.
 
@@ -249,7 +239,9 @@ def _inverse(self, z: Tensor, conditions: Tensor = None, **kwargs) -> Tensor:
  z : Tensor
  Samples from a standard normal distribution
  conditions : Tensor, optional, default: None
- Conditions for a approximate conditional distribution
+ Conditions for the approximate conditional distribution
+ training : bool, optional, default: True
+ Whether internal layers (e.g., dropout) should behave in train or inference mode.
  **kwargs : dict, optional, default: {}
  Additional keyword arguments. Include `steps` (default: 10) to
  adjust the number of sampling steps.
@@ -263,15 +255,17 @@ def _inverse(self, z: Tensor, conditions: Tensor = None, **kwargs) -> Tensor:
  x = keras.ops.copy(z) * self.max_time
  discretized_time = keras.ops.flip(self._discretize_time(steps), axis=-1)
  t = keras.ops.full((*keras.ops.shape(x)[:-1], 1), discretized_time[0], dtype=x.dtype)
- x = self.consistency_function(x, t, conditions=conditions)
+
+ x = self.consistency_function(x, t, conditions=conditions, training=training)
+
  for n in range(1, steps):
  noise = keras.random.normal(keras.ops.shape(x), dtype=keras.ops.dtype(x), seed=self.seed_generator)
  x_n = x + keras.ops.sqrt(keras.ops.square(discretized_time[n]) - self.eps**2) * noise
  t = keras.ops.full_like(t, discretized_time[n])
- x = self.consistency_function(x_n, t, conditions=conditions)
+ x = self.consistency_function(x_n, t, conditions=conditions, training=training)
  return x
 
- def consistency_function(self, x: Tensor, t: Tensor, conditions: Tensor = None, **kwargs) -> Tensor:
+ def consistency_function(self, x: Tensor, t: Tensor, conditions: Tensor = None, training: bool = False) -> Tensor:
  """Compute consistency function.
 
  Parameters
@@ -282,16 +276,16 @@ def consistency_function(self, x: Tensor, t: Tensor, conditions: Tensor = None,
  Vector of time samples in [eps, T]
  conditions : Tensor
  The conditioning vector
- **kwargs : dict, optional, default: {}
- Additional keyword arguments passed to the network.
+ training : bool, optional, default: True
+ Whether internal layers (e.g., dropout) should behave in train or inference mode.
  """
 
  if conditions is not None:
  xtc = ops.concatenate([x, t, conditions], axis=-1)
  else:
  xtc = ops.concatenate([x, t], axis=-1)
 
- f = self.output_projector(self.subnet(xtc, **kwargs))
+ f = self.output_projector(self.subnet(xtc, training=training))
 
  # Compute skip and out parts (vectorized, since self.sigma2 is of shape (1, input_dim)
  # Thus, we can do a cross product with the time vector which is (batch_size, 1) for

bayesflow/networks/coupling_flow/coupling_flow.py

@@ -152,7 +152,7 @@ def _forward(
  z = x
  log_det = keras.ops.zeros(keras.ops.shape(x)[:-1])
  for layer in self.invertible_layers:
- z, det = layer(z, conditions=conditions, inverse=False, training=training, **kwargs)
+ z, det = layer(z, conditions=conditions, inverse=False, training=training)
  log_det += det
 
  if density:
@@ -168,7 +168,7 @@ def _inverse(
  x = z
  log_det = keras.ops.zeros(keras.ops.shape(z)[:-1])
  for layer in reversed(self.invertible_layers):
- x, det = layer(x, conditions=conditions, inverse=True, training=training, **kwargs)
+ x, det = layer(x, conditions=conditions, inverse=True, training=training)
  log_det += det
 
  if density: