Deal with inference_network.log_prob to return dict (as PointInferenceNetwork does)

Author: han-ol

bayesflow/approximators/continuous_approximator.py

@@ -458,7 +458,7 @@ def log_prob(self, data: Mapping[str, np.ndarray], **kwargs) -> np.ndarray | dic
         # change of variables formula
         log_det_jac = log_det_jac.get("inference_variables")
         if log_det_jac is not None:
-            log_prob = log_prob + log_det_jac
+            log_prob = keras.tree.map_structure(lambda x: x + log_det_jac, log_prob)
 
         return log_prob