Merge remote-tracking branch 'upstream/dev' into feat-test-compatibility

Author: vpratz

bayesflow/adapters/adapter.py

@@ -820,7 +820,7 @@ def split(self, key: str, *, into: Sequence[str], indices_or_sections: int | Seq
 
         return self
 
-    def squeeze(self, keys: str | Sequence[str], *, axis: int | tuple):
+    def squeeze(self, keys: str | Sequence[str], *, axis: int | Sequence[int]):
         """Append a :py:class:`~transforms.Squeeze` transform to the adapter.
 
         Parameters

bayesflow/adapters/transforms/nan_to_num.py

@@ -13,12 +13,14 @@ class NanToNum(Transform):
 
     Parameters
     ----------
-    default_value : float
-        Value to substitute wherever data is NaN.
-    return_mask : bool, default=False
-        If True, a mask array will be returned under a new key.
-    mask_prefix : str, default='mask_'
-        Prefix for the mask key in the output dictionary.
+    key : str
+        The variable key to look for in the simulation data dict.
+    default_value : float, optional
+        Value to substitute wherever data is NaN. Default is 0.0.
+    return_mask : bool, optional
+        If True, a mask array will be returned under a new key. Default is False.
+    mask_prefix : str, optional
+        Prefix for the mask key in the output dictionary. Default is 'mask_'.
     """
 
     def __init__(self, key: str, default_value: float = 0.0, return_mask: bool = False, mask_prefix: str = "mask"):
@@ -81,10 +83,10 @@ def inverse(self, data: dict[str, any], **kwargs) -> dict[str, any]:
         values = data[self.key]
 
         if not self.return_mask:
-            values[values == self.default_value] = np.nan  # we assume default_value is not in data
+            # assumes default_value is not in nan
+            values[values == self.default_value] = np.nan
         else:
             mask_array = data[self.mask_key].astype(bool)
-            # Put NaNs where mask is 0
             values[~mask_array] = np.nan
 
         data[self.key] = values

bayesflow/adapters/transforms/nnpe.py

@@ -65,8 +65,8 @@ class NNPE(ElementwiseTransform):
     def __init__(
         self,
         *,
-        spike_scale: float | np.ndarray | None = None,
-        slab_scale: float | np.ndarray | None = None,
+        spike_scale: np.typing.ArrayLike | None = None,
+        slab_scale: np.typing.ArrayLike | None = None,
         per_dimension: bool = True,
         seed: int | None = None,
     ):
@@ -80,14 +80,14 @@ def __init__(
     def _resolve_scale(
         self,
         name: str,
-        passed: float | np.ndarray | None,
+        passed: np.typing.ArrayLike | None,
         default: float,
         data: np.ndarray,
     ) -> np.ndarray | float:
         """
         Determine spike/slab scale:
-         - If passed is None: Automatic determination via default * std(data) (per‐dimension or global).
-         - Else: validate & cast passed to the correct shape/type.
+         - If `passed` is None: Automatic determination via default * std(data) (per‐dimension or global).
+         - Else: Validate & cast `passed` to the correct shape/type.
 
         Parameters
         ----------
@@ -103,8 +103,8 @@ def _resolve_scale(
 
         Returns
         -------
-        float or np.ndarray
-            The resolved scale, either as a scalar (if per_dimension=False) or an 1D array of length data.shape[-1]
+        np.ndarray
+            The resolved scale, either as a 0D array (if per_dimension=False) or an 1D array of length data.shape[-1]
             (if per_dimension=True).
         """
 
@@ -119,22 +119,22 @@ def _resolve_scale(
 
         # If no scale is passed, determine scale automatically given the dimensionwise or global std
         if passed is None:
-            return default * std
+            return np.array(default * std)
         # If a scale is passed, check if the passed shape matches the expected shape
         else:
-            if self.per_dimension:
+            try:
                 arr = np.asarray(passed, dtype=float)
-                if arr.shape != expected_shape or arr.ndim != 1:
+            except Exception as e:
+                raise TypeError(f"{name}: expected values convertible to float, got {type(passed).__name__}") from e
+
+            if self.per_dimension:
+                if arr.ndim != 1 or arr.shape != expected_shape:
                     raise ValueError(f"{name}: expected array of shape {expected_shape}, got {arr.shape}")
                 return arr
             else:
-                try:
-                    scalar = float(passed)
-                except TypeError:
-                    raise TypeError(f"{name}: expected a scalar convertible to float, got type {type(passed).__name__}")
-                except ValueError:
-                    raise ValueError(f"{name}: expected a scalar convertible to float, got value {passed!r}")
-                return scalar
+                if arr.ndim != 0:
+                    raise ValueError(f"{name}: expected scalar, got array of shape {arr.shape}")
+                return arr
 
     def forward(self, data: np.ndarray, stage: str = "inference", **kwargs) -> np.ndarray:
         """
@@ -173,7 +173,7 @@ def forward(self, data: np.ndarray, stage: str = "inference", **kwargs) -> np.nd
         return data + noise
 
     def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
-        """Non-invertible transform."""
+        # Non-invertible transform
         return data
 
     def get_config(self) -> dict:

bayesflow/adapters/transforms/squeeze.py

@@ -1,5 +1,6 @@
 import numpy as np
 
+from collections.abc import Sequence
 from bayesflow.utils.serialization import serializable, serialize
 
 from .elementwise_transform import ElementwiseTransform
@@ -29,8 +30,10 @@ class Squeeze(ElementwiseTransform):
     It is recommended to precede this transform with a :class:`~bayesflow.adapters.transforms.ToArray` transform.
     """
 
-    def __init__(self, *, axis: int | tuple):
+    def __init__(self, *, axis: int | Sequence[int]):
         super().__init__()
+        if isinstance(axis, Sequence):
+            axis = tuple(axis)
         self.axis = axis
 
     def get_config(self) -> dict:

bayesflow/approximators/backend_approximators/backend_approximator.py

@@ -6,8 +6,6 @@
 match keras.backend.backend():
     case "jax":
         from .jax_approximator import JAXApproximator as BaseBackendApproximator
-    case "numpy":
-        from .numpy_approximator import NumpyApproximator as BaseBackendApproximator
     case "tensorflow":
         from .tensorflow_approximator import TensorFlowApproximator as BaseBackendApproximator
     case "torch":

bayesflow/approximators/backend_approximators/jax_approximator.py

@@ -56,6 +56,20 @@ def stateless_compute_metrics(
         variables and returns both the loss and auxiliary information for
         further updates.
 
+        Things we do for specifically jax:
+
+        1. Accept trainable variables as the first argument
+            (can be at any position as indicated by the argnum parameter
+             in autograd, but needs to be an explicit arg)
+        2. Accept, potentially modify, and return other state variables
+        3. Return just the loss tensor as the first value
+        4. Return all other values in a tuple as the second value
+
+        This ensures:
+
+        1. The function is stateless
+        2. The function can be differentiated with jax autograd
+
         Parameters
         ----------
         trainable_variables : Any

bayesflow/approximators/backend_approximators/numpy_approximator.py

@@ -3,6 +3,7 @@
 import numpy as np
 
 import keras
+import warnings
 
 from bayesflow.adapters import Adapter
 from bayesflow.networks import InferenceNetwork, SummaryNetwork
@@ -98,8 +99,8 @@ def build(self, data_shapes: dict[str, tuple[int] | dict[str, dict]]) -> None:
             ]
             self.standardize_layers = {var: Standardization(trainable=False) for var in self.standardize}
 
-        # Build all standardization layers, if present
-        for var, layer in getattr(self, "standardize_layers", {}).items():
+        # Build all standardization layers
+        for var, layer in self.standardize_layers.items():
             layer.build(data_shapes[var])
 
         self.built = True
@@ -448,7 +449,7 @@ def sample(
         conditions = self._prepare_data(conditions, **kwargs)
 
         # Remove any superfluous keys, just retain actual conditions
-        conditions = {k: v for k, v in conditions.items() if k in ContinuousApproximator.CONDITION_KEYS}
+        conditions = {k: v for k, v in conditions.items() if k in self.CONDITION_KEYS}
 
         # Sample and undo optional standardization
         samples = self._sample(num_samples=num_samples, **conditions, **kwargs)
@@ -485,7 +486,7 @@ def _prepare_data(
             ldj_inference = None
 
         # Standardize conditions
-        for key in ContinuousApproximator.CONDITION_KEYS:
+        for key in self.CONDITION_KEYS:
             if key in self.standardize and key in data:
                 data[key] = self.standardize_layers[key](data[key])
 
@@ -514,8 +515,12 @@ def _sample(
         summary_variables: Tensor = None,
         **kwargs,
     ) -> Tensor:
-        if (self.summary_network is None) != (summary_variables is None):
-            raise ValueError("Summary variables and summary network must be used together.")
+        if self.summary_network is None:
+            if summary_variables is not None:
+                raise ValueError("Cannot use summary variables without a summary network.")
+        else:
+            if summary_variables is None:
+                raise ValueError("Summary variables are required when a summary network is present.")
 
         if self.summary_network is not None:
             summary_outputs = self.summary_network(
@@ -539,7 +544,7 @@ def _sample(
             batch_shape, conditions=inference_conditions, **filter_kwargs(kwargs, self.inference_network.sample)
         )
 
-    def summaries(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray:
+    def summarize(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray:
         """
         Computes the learned summary statistics of given summary variables.
 
@@ -570,6 +575,14 @@ def summaries(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray:
 
         return summaries
 
+    def summaries(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray:
+        """
+        .. deprecated:: 2.0.4
+            `summaries` will be removed in version 2.0.5, it was renamed to `summarize` which should be used instead.
+        """
+        warnings.warn("`summaries` was renamed to `summarize` and will be removed in version 2.0.5.", FutureWarning)
+        return self.summarize(data=data, **kwargs)
+
     def log_prob(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray:
         """
         Computes the log-probability of given data under the model. The `data` dictionary is preprocessed using the
@@ -606,8 +619,12 @@ def _log_prob(
         summary_variables: Tensor = None,
         **kwargs,
     ) -> Tensor:
-        if (self.summary_network is None) != (summary_variables is None):
-            raise ValueError("Summary variables and summary network must be used together.")
+        if self.summary_network is None:
+            if summary_variables is not None:
+                raise ValueError("Cannot use summary variables without a summary network.")
+        else:
+            if summary_variables is None:
+                raise ValueError("Summary variables are required when a summary network is present.")
 
         if self.summary_network is not None:
             summary_outputs = self.summary_network(

bayesflow/approximators/continuous_approximator.py

@@ -2,6 +2,7 @@
 
 import keras
 import numpy as np
+import warnings
 
 from bayesflow.adapters import Adapter
 from bayesflow.datasets import OnlineDataset
@@ -92,11 +93,11 @@ def build(self, data_shapes: dict[str, tuple[int] | dict[str, dict]]) -> None:
 
         # Set up standardization layers if requested
         if self.standardize == "all":
-            self.standardize = [var for var in ModelComparisonApproximator.CONDITION_KEYS if var in data_shapes]
+            self.standardize = [var for var in self.CONDITION_KEYS if var in data_shapes]
             self.standardize_layers = {var: Standardization(trainable=False) for var in self.standardize}
 
         # Build all standardization layers
-        for var, layer in getattr(self, "standardize_layers", {}).items():
+        for var, layer in self.standardize_layers.items():
             layer.build(data_shapes[var])
 
         self.built = True
@@ -242,7 +243,7 @@ def compute_metrics(
     def fit(
         self,
         *,
-        adapter: Adapter | str = "auto",
+        adapter: Adapter = "auto",
         dataset: keras.utils.PyDataset = None,
         simulator: ModelComparisonSimulator = None,
         simulators: Sequence[Simulator] = None,
@@ -256,7 +257,7 @@ def fit(
 
         Parameters
         ----------
-        adapter : Adapter or str, optional
+        adapter : Adapter or 'auto', optional
             The data adapter that will make the simulated / real outputs neural-network friendly.
         dataset : keras.utils.PyDataset, optional
             A dataset containing simulations for training. If provided, `simulator` must be None.
@@ -392,19 +393,22 @@ def predict(
         conditions = self.adapter(conditions, strict=False, stage="inference", **kwargs)
 
         # 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.CONDITION_KEYS}
+        conditions = {k: v for k, v in conditions.items() if k in self.CONDITION_KEYS}
         conditions = keras.tree.map_structure(keras.ops.convert_to_tensor, conditions)
 
         # Optionally standardize conditions
-        for key in ModelComparisonApproximator.CONDITION_KEYS:
+        for key in self.CONDITION_KEYS:
             if key in conditions and key in self.standardize:
                 conditions[key] = self.standardize_layers[key](conditions[key])
 
         output = self._predict(**conditions, **kwargs)
 
-        return keras.ops.convert_to_numpy(keras.ops.softmax(output) if probs else output)
+        if probs:
+            output = keras.ops.softmax(output)
 
-    def summaries(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray:
+        return keras.ops.convert_to_numpy(output)
+
+    def summarize(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray:
         """
         Computes the learned summary statistics of given summary variables.
 
@@ -435,6 +439,14 @@ def summaries(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray:
 
         return summaries
 
+    def summaries(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray:
+        """
+        .. deprecated:: 2.0.4
+            `summaries` will be removed in version 2.0.5, it was renamed to `summarize` which should be used instead.
+        """
+        warnings.warn("`summaries` was renamed to `summarize` and will be removed in version 2.0.5.", FutureWarning)
+        return self.summarize(data=data, **kwargs)
+
     def _compute_logits(self, classifier_conditions: Tensor) -> Tensor:
         """Helper to compute projected logits from the classifier network."""
         logits = self.classifier_network(classifier_conditions)