Update docs and typehints

Author: stefanradev93

bayesflow/networks/inference_network.py

@@ -1,6 +1,8 @@
-import keras
+from typing import Literal
 from collections.abc import Sequence
 
+import keras
+
 from bayesflow.types import Shape, Tensor
 from bayesflow.utils import layer_kwargs, find_distribution
 from bayesflow.utils.decorators import allow_batch_size
@@ -9,7 +11,28 @@
 
 @serializable("bayesflow.networks")
 class InferenceNetwork(keras.Layer):
-    def __init__(self, base_distribution: str = "normal", *, metrics: Sequence[keras.Metric] | None = None, **kwargs):
+    def __init__(
+        self,
+        base_distribution: Literal["normal", "student", "mixture"] | keras.Layer = "normal",
+        *,
+        metrics: Sequence[keras.Metric] | None = None,
+        **kwargs,
+    ):
+        """
+        Constructs an inference network using a specified base distribution and optional custom metrics.
+        Use this interface for custom inference networks.
+
+        Parameters
+        ----------
+        base_distribution : Literal["normal", "student", "mixture"] or keras.Layer
+            Name or the actual base distribution to use. Passed to `find_distribution` to
+            obtain the corresponding distribution object.
+        metrics : Sequence[keras.Metric] or None, optional
+            Sequence of custom Keras Metric instances to compute during training
+            and evaluation. If `None`, no custom metrics are used.
+        **kwargs
+            Additional keyword arguments forwarded to the `keras.Layer` constructor.
+        """
         self.custom_metrics = metrics
         super().__init__(**layer_kwargs(kwargs))
         self.base_distribution = find_distribution(base_distribution)
@@ -70,7 +93,7 @@ def compute_metrics(
 
         if stage != "training" and any(self.metrics):
             # compute sample-based metrics
-            samples = self.sample((keras.ops.shape(x)[0],), conditions=conditions)
+            samples = self.sample(batch_shape=(keras.ops.shape(x)[0],), conditions=conditions)
 
             for metric in self.metrics:
                 metrics[metric.name] = metric(samples, x)

bayesflow/networks/mlp/mlp.py

@@ -55,9 +55,11 @@ def __init__(
         dropout : float or None, optional
             Dropout rate applied within the MLP layers for regularization. Default is 0.05.
         norm: str, optional
-
+            Type of learnable normalization to be used (e.g., "batch" or "layer"). Default is None.
         spectral_normalization : bool, optional
             Whether to apply spectral normalization to stabilize training. Default is False.
+        metrics: Sequence[keras.Metric], optional
+            A sequence of callable metrics following keras' `Metric` signature. Default is None.
         **kwargs
             Additional keyword arguments passed to the Keras layer initialization.
         """

bayesflow/networks/summary_network.py

@@ -11,6 +11,29 @@
 @serializable("bayesflow.networks")
 class SummaryNetwork(keras.Layer):
     def __init__(self, base_distribution: str = None, *, metrics: Sequence[keras.Metric] | None = None, **kwargs):
+        """
+        Builds a summary network with an optional base distribution and custom metrics. Use this class
+        as an interface for custom summary networks.
+
+        Important: If a base distribution is passed, the summary outputs will be optimized to follow
+        said distribution, as described in [1].
+
+        [1] Schmitt, M., Bürkner, P. C., Köthe, U., & Radev, S. T. (2023).
+        Detecting model misspecification in amortized Bayesian inference with neural networks.
+        In DAGM German Conference on Pattern Recognition (pp. 541-557). Cham: Springer Nature Switzerland.
+
+        Parameters
+        ----------
+        base_distribution : str or None, default None
+            Name of the base distribution to use. If `None`, a default distribution
+            is chosen. Passed to `find_distribution` to obtain the corresponding
+            distribution object.
+        metrics : Sequence[keras.Metric] or None, optional
+            Sequence of custom Keras Metric instances to compute during training
+            and evaluation. If `None`, no custom metrics are used.
+        **kwargs
+            Additional keyword arguments forwarded to the `keras.Layer` constructor.
+        """
         self.custom_metrics = metrics
         super().__init__(**layer_kwargs(kwargs))
         self.base_distribution = find_distribution(base_distribution)