Fix typehints to docs.

Author: stefanradev93

bayesflow/approximators/approximator.py

@@ -1,5 +1,3 @@
-from collections.abc import Mapping
-
 import multiprocessing as mp
 
 import keras
@@ -22,7 +20,7 @@ def build_adapter(cls, **kwargs) -> Adapter:
         # implemented by each respective architecture
         raise NotImplementedError
 
-    def build_from_data(self, data: Mapping[str, any]) -> None:
+    def build_from_data(self, data: dict[str, any]) -> None:
         self.compute_metrics(**filter_kwargs(data, self.compute_metrics), stage="training")
         self.built = True
 
@@ -137,27 +135,3 @@ def fit(self, *, dataset: keras.utils.PyDataset = None, simulator: Simulator = N
             self.build_from_data(mock_data)
 
         return super().fit(dataset=dataset, **kwargs)
-
-    def _batch_size_from_data(self, data: any):
-        """Obtain the batch size from a batch of data.
-
-        To properly weight the metrics for batches of different sizes, the batch size of a given batch of data is
-        required. As the data structure differs between approximators, each concrete approximator has to specify
-        this method.
-
-        Parameters
-        ----------
-        data :
-            The data that are passed to `compute_metrics` as keyword arguments.
-
-        Returns
-        -------
-        batch_size : int
-            The batch size of the given data.
-        """
-        raise NotImplementedError(
-            "Correct calculation of the metrics requires obtaining the batch size from the supplied data "
-            "for proper weighting of metrics for batches with different sizes. Please implement the "
-            "_batch_size_from_data method for your approximator. For a given batch of data, it should "
-            "return the corresponding batch size."
-        )

bayesflow/approximators/backend_approximators/jax_approximator.py

@@ -31,7 +31,7 @@ def compute_metrics(self, *args, **kwargs) -> dict[str, jax.Array]:
         ----------
         *args : tuple
             Positional arguments passed to the metric computation function.
-        **kwargs : dict
+        **kwargs
             Keyword arguments passed to the metric computation function.
 
         Returns
@@ -222,3 +222,28 @@ def _update_metrics(self, loss: jax.Array, metrics_variables: any, sample_weight
         metrics_variables = [scope.get_current_value(v) for v in self.metrics_variables]
 
         return metrics_variables
+
+    # noinspection PyMethodOverriding
+    def _batch_size_from_data(self, data: any) -> int:
+        """Obtain the batch size from a batch of data.
+
+        To properly weigh the metrics for batches of different sizes, the batch size of a given batch of data is
+        required. As the data structure differs between approximators, each concrete approximator has to specify
+        this method.
+
+        Parameters
+        ----------
+        data :
+            The data that are passed to `compute_metrics` as keyword arguments.
+
+        Returns
+        -------
+        batch_size : int
+            The batch size of the given data.
+        """
+        raise NotImplementedError(
+            "Correct calculation of the metrics requires obtaining the batch size from the supplied data "
+            "for proper weighting of metrics for batches with different sizes. Please implement the "
+            "_batch_size_from_data method for your approximator. For a given batch of data, it should "
+            "return the corresponding batch size."
+        )

bayesflow/approximators/backend_approximators/numpy_approximator.py

@@ -27,3 +27,12 @@ def _update_metrics(self, metrics, sample_weight=None):
             except ValueError:
                 self._metrics.append(keras.metrics.Mean(name=name))
                 self._metrics[-1].update_state(value, sample_weight=sample_weight)
+
+    # noinspection PyMethodOverriding
+    def _batch_size_from_data(self, data: any) -> int:
+        raise NotImplementedError(
+            "Correct calculation of the metrics requires obtaining the batch size from the supplied data "
+            "for proper weighting of metrics for batches with different sizes. Please implement the "
+            "_batch_size_from_data method for your approximator. For a given batch of data, it should "
+            "return the corresponding batch size."
+        )

bayesflow/approximators/backend_approximators/tensorflow_approximator.py

@@ -32,7 +32,7 @@ def compute_metrics(self, *args, **kwargs) -> dict[str, tf.Tensor]:
         ----------
         *args : tuple
             Positional arguments passed to the metric computation function.
-        **kwargs : dict
+        **kwargs
             Keyword arguments passed to the metric computation function.
 
         Returns
@@ -114,3 +114,28 @@ def _update_metrics(self, metrics: dict[str, any], sample_weight: tf.Tensor = No
             except ValueError:
                 self._metrics.append(keras.metrics.Mean(name=name))
                 self._metrics[-1].update_state(value, sample_weight=sample_weight)
+
+    # noinspection PyMethodOverriding
+    def _batch_size_from_data(self, data: any) -> int:
+        """Obtain the batch size from a batch of data.
+
+        To properly weigh the metrics for batches of different sizes, the batch size of a given batch of data is
+        required. As the data structure differs between approximators, each concrete approximator has to specify
+        this method.
+
+        Parameters
+        ----------
+        data :
+            The data that are passed to `compute_metrics` as keyword arguments.
+
+        Returns
+        -------
+        batch_size : int
+            The batch size of the given data.
+        """
+        raise NotImplementedError(
+            "Correct calculation of the metrics requires obtaining the batch size from the supplied data "
+            "for proper weighting of metrics for batches with different sizes. Please implement the "
+            "_batch_size_from_data method for your approximator. For a given batch of data, it should "
+            "return the corresponding batch size."
+        )

bayesflow/approximators/backend_approximators/torch_approximator.py

@@ -33,7 +33,7 @@ def compute_metrics(self, *args, **kwargs) -> dict[str, torch.Tensor]:
         ----------
         *args : tuple
             Positional arguments passed to the metric computation function.
-        **kwargs : dict
+        **kwargs
             Keyword arguments passed to the metric computation function.
 
         Returns
@@ -124,3 +124,28 @@ def _update_metrics(self, metrics: dict[str, any], sample_weight: torch.Tensor =
             except ValueError:
                 self._metrics.append(keras.metrics.Mean(name=name))
                 self._metrics[-1].update_state(value, sample_weight=sample_weight)
+
+    # noinspection PyMethodOverriding
+    def _batch_size_from_data(self, data: any) -> int:
+        """Obtain the batch size from a batch of data.
+
+        To properly weigh the metrics for batches of different sizes, the batch size of a given batch of data is
+        required. As the data structure differs between approximators, each concrete approximator has to specify
+        this method.
+
+        Parameters
+        ----------
+        data :
+            The data that are passed to `compute_metrics` as keyword arguments.
+
+        Returns
+        -------
+        batch_size : int
+            The batch size of the given data.
+        """
+        raise NotImplementedError(
+            "Correct calculation of the metrics requires obtaining the batch size from the supplied data "
+            "for proper weighting of metrics for batches with different sizes. Please implement the "
+            "_batch_size_from_data method for your approximator. For a given batch of data, it should "
+            "return the corresponding batch size."
+        )