Bayes Estimators: Lossfunctions with flexible signature

Author: han-ol

bayesflow/amortizers.py

@@ -31,7 +31,7 @@
 from bayesflow.default_settings import DEFAULT_KEYS
 from bayesflow.exceptions import ConfigurationError, SummaryStatsError
 from bayesflow.helper_functions import check_tensor_sanity
-from bayesflow.losses import log_loss, mmd_summary_space
+from bayesflow.losses import log_loss, mmd_summary_space, norm_diff
 from bayesflow.networks import EvidentialNetwork
 
 
@@ -1337,7 +1337,7 @@ def compute_loss(self, input_dict, **kwargs):
         """
 
         net_out = self(input_dict, **kwargs)
-        loss = tf.reduce_mean(self.loss_fn(net_out - input_dict[DEFAULT_KEYS["parameters"]]))
+        loss = tf.reduce_mean(self.loss_fn(net_out, input_dict[DEFAULT_KEYS["parameters"]]))
         return loss
 
     def _compute_summary_condition(self, summary_conditions, direct_conditions, **kwargs):
@@ -1366,4 +1366,4 @@ def _determine_loss(self, loss_fun, norm_ord):
         # In case of user-provided loss, override norm order
         if loss_fun is not None:
             return loss_fun
-        return partial(tf.norm, ord=norm_ord, axis=-1)
+        return partial(norm_diff, ord=norm_ord, axis=-1)

bayesflow/losses.py

@@ -184,3 +184,20 @@ def log_loss(model_indices, preds, evidential=False, label_smoothing=0.01):
     # Actual loss + regularization (if given)
     loss = -tf.reduce_mean(tf.reduce_sum(model_indices * tf.math.log(preds), axis=1))
     return loss
+
+
+def norm_diff(tensor_a, tensor_b, axis=None, ord='euclidean'):
+    """
+    Wrapper around tf.norm that computes the norm of the difference between two tensors along the specified axis.
+
+    Parameters
+    ----------
+    tensor_a : A Tensor.
+    tensor_b : A Tensor. Must be the same shape as tensor_a.
+    axis     : Any or None
+        Axis along which to compute the norm of the difference. Default is None.
+    ord      : int or str
+        Order of the norm. Supports 'euclidean' and other norms supported by tf.norm. Default is 'euclidean'.
+    """
+    difference = tensor_a - tensor_b
+    return tf.norm(difference, ord=ord, axis=axis)