Add small doc improvements

Author: stefanradev93

bayesflow/approximators/approximator.py

@@ -142,7 +142,8 @@ 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 approximator has to specify this method.
+        required. As the data structure differs between approximators, each concrete approximator has to specify
+        this method.
 
         Parameters
         ----------

bayesflow/approximators/continuous_approximator.py

@@ -492,5 +492,9 @@ def _log_prob(
             **filter_kwargs(kwargs, self.inference_network.log_prob),
         )
 
-    def _batch_size_from_data(self, data: Mapping[str, any]):
+    def _batch_size_from_data(self, data: Mapping[str, any]) -> int:
+        """
+        Fetches the current batch size from an input dictionary. Can only be used during training when
+        inference variables as present.
+        """
         return keras.ops.shape(data["inference_variables"])[0]

bayesflow/approximators/model_comparison_approximator.py

@@ -379,5 +379,9 @@ def summaries(self, data: Mapping[str, np.ndarray], **kwargs):
         summaries = self.summary_network(summary_variables, **filter_kwargs(kwargs, self.summary_network.call))
         return summaries
 
-    def _batch_size_from_data(self, data: Mapping[str, any]):
+    def _batch_size_from_data(self, data: Mapping[str, any]) -> int:
+        """
+        Fetches the current batch size from an input dictionary. Can only be used during training when
+        model indices as present.
+        """
         return keras.ops.shape(data["model_indices"])[0]