Update dispatching distributions

Author: stefanradev93

bayesflow/networks/inference_network.py

@@ -1,8 +1,7 @@
 import keras
 
-from bayesflow.distributions import find_distribution
 from bayesflow.types import Shape, Tensor
-from bayesflow.utils import layer_kwargs
+from bayesflow.utils import layer_kwargs, find_distribution
 from bayesflow.utils.decorators import allow_batch_size
 
 

bayesflow/networks/summary_network.py

@@ -1,9 +1,8 @@
 import keras
 
-from bayesflow.distributions import find_distribution
 from bayesflow.metrics.functional import maximum_mean_discrepancy
 from bayesflow.types import Tensor
-from bayesflow.utils import layer_kwargs
+from bayesflow.utils import layer_kwargs, find_distribution
 from bayesflow.utils.decorators import sanitize_input_shape
 from bayesflow.utils.serialization import deserialize
 

bayesflow/utils/__init__.py

@@ -7,8 +7,11 @@
  logging,
  numpy_utils,
 )
+
 from .callbacks import detailed_loss_callback
+
 from .devices import devices
+
 from .dict_utils import (
  convert_args,
  convert_kwargs,
@@ -20,30 +23,48 @@
  split_arrays,
  squeeze_inner_estimates_dict,
 )
-from .dispatch import find_network, find_permutation, find_pooling, find_recurrent_net
+
+from .dispatch import (
+ find_network,
+ find_permutation,
+ find_pooling,
+ find_recurrent_net,
+ find_summary_network,
+ find_inference_network,
+ find_distribution,
+)
+
 from .ecdf import simultaneous_ecdf_bands, ranks
+
 from .functional import batched_call
+
 from .git import (
  issue_url,
  pull_url,
  repo_url,
 )
+
 from .hparam_utils import find_batch_size, find_memory_budget
+
 from .integrate import (
  integrate,
 )
+
 from .io import (
  pickle_load,
  format_bytes,
  parse_bytes,
 )
+
 from .jacobian import (
  jacobian,
  jacobian_trace,
  jvp,
  vjp,
 )
+
 from .optimal_transport import optimal_transport
+
 from .plot_utils import (
  check_estimates_prior_shapes,
  prepare_plot_data,
@@ -53,6 +74,7 @@
  add_metric,
 )
 from .serialization import serialize_value_or_type, deserialize_value_or_type
+
 from .tensor_utils import (
  concatenate_valid,
  expand,
@@ -75,9 +97,10 @@
  fill_triangular_matrix,
  weighted_mean,
 )
+
 from .classification import calibration_curve, confusion_matrix
+
 from .validators import check_lengths_same
-from .workflow_utils import find_inference_network, find_summary_network
 
 from ._docs import _add_imports_to_all