fix order of prior scores

Author: arrjon

bayesflow/approximators/continuous_approximator.py

@@ -712,7 +712,7 @@ def compute_prior_score_pre(_samples: Tensor) -> Tensor:
                     prior_score[key] -= expand_right_as(log_det_jac[key], prior_score[key])
 
             prior_score = keras.tree.map_structure(keras.ops.convert_to_tensor, prior_score)
-            out = keras.ops.concatenate(list(prior_score.values()), axis=-1)
+            out = keras.ops.concatenate([prior_score[key] for key in adapted_samples], axis=-1)
             return out - keras.ops.expand_dims(log_det_jac_standardize, axis=-1)
 
         # Test prior score function, useful for debugging

tests/test_approximators/conftest.py

@@ -261,7 +261,7 @@ def transforming_adapter():
     # Apply log transformation to scale parameters (to make them unbounded)
     adapter.log(["scale"])
 
-    adapter.concatenate(["loc", "scale"], into="inference_variables")
+    adapter.concatenate(["scale", "loc"], into="inference_variables")
     adapter.concatenate(["conditions"], into="inference_conditions")
     adapter.keep(["inference_variables", "inference_conditions"])
     return adapter

tests/test_approximators/test_compositional_prior_score.py

@@ -88,7 +88,7 @@ def compute_prior_score_pre(_samples):
                     prior_score[key] -= expand_right_as(log_det_jac[key], prior_score[key])
 
             prior_score = keras.tree.map_structure(keras.ops.convert_to_tensor, prior_score)
-            out = keras.ops.concatenate(list(prior_score.values()), axis=-1)
+            out = keras.ops.concatenate([prior_score[key] for key in adapted_samples], axis=-1)
             return out - keras.ops.expand_dims(log_det_jac_standardize, axis=-1)
 
         return compute_prior_score_pre
@@ -102,8 +102,9 @@ def compute_prior_score_pre(_samples):
 
     # With Jacobian correction: score_transformed = score_original - log|J|
     old_scores = mock_prior_score_original_space(dummy_data_dict)
-    det_jac_scale = y_samples[0, 2:].sum()
-    expected_scores = np.array([old_scores["loc"][0], old_scores["scale"][0] - det_jac_scale]).flatten()
+    # order of parameters is flipped due to concatenation in adapter
+    det_jac_scale = y_samples[0, :2].sum()
+    expected_scores = np.array([old_scores["scale"][0] - det_jac_scale, old_scores["loc"][0]]).flatten()
 
     # Check that scores are reasonably close
     np.testing.assert_allclose(scores_np, expected_scores, rtol=1e-5, atol=1e-6)