Add split arrays as an option to sample

Author: stefanradev93

bayesflow/approximators/continuous_approximator.py

@@ -11,7 +11,7 @@
 from bayesflow.adapters import Adapter
 from bayesflow.networks import InferenceNetwork, SummaryNetwork
 from bayesflow.types import Tensor
-from bayesflow.utils import logging
+from bayesflow.utils import logging, split_arrays
 from .approximator import Approximator
 
 
@@ -136,6 +136,7 @@ def sample(
         *,
         num_samples: int,
         conditions: dict[str, np.ndarray],
+        split: bool = False,
         **kwargs,
     ) -> dict[str, np.ndarray]:
         conditions = self.adapter(conditions, strict=False, stage="inference", **kwargs)
@@ -144,6 +145,8 @@ def sample(
         conditions = keras.tree.map_structure(keras.ops.convert_to_numpy, conditions)
         conditions = self.adapter(conditions, inverse=True, strict=False, **kwargs)
 
+        if split:
+            conditions = split_arrays(conditions, axis=-1)
         return conditions
 
     def _sample(

bayesflow/utils/__init__.py

@@ -3,13 +3,7 @@
     logging,
     numpy_utils,
 )
-from .dict_utils import (
-    convert_args,
-    convert_kwargs,
-    filter_kwargs,
-    keras_kwargs,
-    split_tensors,
-)
+from .dict_utils import convert_args, convert_kwargs, filter_kwargs, keras_kwargs, split_tensors, split_arrays
 from .dispatch import find_distribution, find_network, find_permutation, find_pooling, find_recurrent_net
 from .ecdf import simultaneous_ecdf_bands, ranks
 from .functional import batched_call

bayesflow/utils/dict_utils.py

@@ -105,6 +105,24 @@ def split_tensors(data: Mapping[any, Tensor], axis: int = -1) -> Mapping[any, Te
     return result
 
 
+def split_arrays(data: Mapping[any, np.ndarray], axis: int = -1) -> Mapping[any, np.ndarray]:
+    """Split tensors in the dictionary along the given axis."""
+    result = {}
+
+    for key, value in data.items():
+        if value.shape[axis] == 1:
+            result[key] = np.squeeze(value, axis=axis)
+            continue
+
+        splits = np.split(value, value.shape[axis], axis=axis)
+        splits = [np.squeeze(split, axis=axis) for split in splits]
+
+        for i, split in enumerate(splits):
+            result[f"{key}_{i + 1}"] = split
+
+    return result
+
+
 def dicts_to_arrays(
     post_variables: dict[str, np.ndarray] | np.ndarray,
     prior_variables: dict[str, np.ndarray] | np.ndarray = None,