Move part of nested dictionary operations to dict_utils

Author: han-ol

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 filter_kwargs, logging, split_arrays
+from bayesflow.utils import filter_kwargs, logging, split_arrays, squeeze_inner_estimates_dict
 from .approximator import Approximator
 
 
@@ -232,16 +232,10 @@ def estimate(
             for variable_name in samples.keys()
         }
 
-        def squeeze_dict(d):
-            if len(d.keys()) == 1 and "value" in d.keys():
-                return d["value"]
-            else:
-                return d
-
         # remove unnecessary nesting
         estimates = {
             variable_name: {
-                outer_key: squeeze_dict(estimates[variable_name][outer_key])
+                outer_key: squeeze_inner_estimates_dict(estimates[variable_name][outer_key])
                 for outer_key in estimates[variable_name].keys()
             }
             for variable_name in estimates.keys()

bayesflow/approximators/point_approximator.py

@@ -5,7 +5,7 @@
 )
 
 from bayesflow.types import Tensor
-from bayesflow.utils import filter_kwargs, split_arrays
+from bayesflow.utils import filter_kwargs, split_arrays, squeeze_inner_estimates_dict
 from .continuous_approximator import ContinuousApproximator
 
 
@@ -62,16 +62,10 @@ def estimate(
             for variable_name in inference_variable_names
         }
 
-        def squeeze_dict(d):
-            if len(d.keys()) == 1 and "value" in d.keys():
-                return d["value"]
-            else:
-                return d
-
         # remove unnecessary nesting
         conditions = {
             variable_name: {
-                outer_key: squeeze_dict(conditions[variable_name][outer_key])
+                outer_key: squeeze_inner_estimates_dict(conditions[variable_name][outer_key])
                 for outer_key in conditions[variable_name].keys()
             }
             for variable_name in conditions.keys()

bayesflow/utils/__init__.py

@@ -12,6 +12,7 @@
     keras_kwargs,
     split_tensors,
     split_arrays,
+    squeeze_inner_estimates_dict,
 )
 from .dispatch import find_distribution, find_network, find_permutation, find_pooling, find_recurrent_net
 from .ecdf import simultaneous_ecdf_bands, ranks

bayesflow/utils/dict_utils.py

@@ -318,3 +318,15 @@ def dicts_to_arrays(
         targets=targets,
         priors=priors,
     )
+
+
+def squeeze_inner_estimates_dict(estimates):
+    """If a dictionary has only one key-value pair and the key is "value", return only its value.
+    Otherwise, return the unchanged dictionary.
+
+    This method helps to remove unnecessary nesting levels.
+    """
+    if len(estimates.keys()) == 1 and "value" in estimates.keys():
+        return estimates["value"]
+    else:
+        return estimates