Remove skip instructions across approximator tests regarding MVNScore

Author: han-ol

tests/test_approximators/test_fit.py

@@ -3,7 +3,6 @@
 import pytest
 import io
 from contextlib import redirect_stdout
-from tests.utils import check_approximator_multivariate_normal_score
 
 
 @pytest.mark.skip(reason="not implemented")
@@ -20,9 +19,6 @@ def test_fit(amortizer, dataset):
 
 
 def test_loss_progress(approximator, train_dataset, validation_dataset):
-    # as long as MultivariateNormalScore is unstable, skip fit progress test
-    check_approximator_multivariate_normal_score(approximator)
-
     approximator.compile(optimizer="AdamW")
     num_epochs = 3
 

tests/test_approximators/test_log_prob.py

@@ -1,12 +1,10 @@
 import keras
 import numpy as np
-from tests.utils import check_combination_simulator_adapter, check_approximator_multivariate_normal_score
+from tests.utils import check_combination_simulator_adapter
 
 
 def test_approximator_log_prob(approximator, simulator, batch_size, adapter):
     check_combination_simulator_adapter(simulator, adapter)
-    # as long as MultivariateNormalScore is unstable, skip
-    check_approximator_multivariate_normal_score(approximator)
 
     num_batches = 4
     data = simulator.sample((num_batches * batch_size,))

tests/test_approximators/test_sample.py

@@ -1,12 +1,10 @@
 import numpy as np
 import keras
-from tests.utils import check_combination_simulator_adapter, check_approximator_multivariate_normal_score
+from tests.utils import check_combination_simulator_adapter
 
 
 def test_approximator_sample(approximator, simulator, batch_size, adapter):
     check_combination_simulator_adapter(simulator, adapter)
-    # as long as MultivariateNormalScore is unstable, skip
-    check_approximator_multivariate_normal_score(approximator)
 
     num_batches = 4
     data = simulator.sample((num_batches * batch_size,))
@@ -23,8 +21,6 @@ def test_approximator_sample(approximator, simulator, batch_size, adapter):
 
 def test_approximator_sample_keep_conditions(approximator, simulator, batch_size, adapter):
     check_combination_simulator_adapter(simulator, adapter)
-    # as long as MultivariateNormalScore is unstable, skip
-    check_approximator_multivariate_normal_score(approximator)
 
     num_batches = 4
     data = simulator.sample((num_batches * batch_size,))
@@ -39,14 +35,18 @@ def test_approximator_sample_keep_conditions(approximator, simulator, batch_size
 
     assert isinstance(samples_and_conditions, dict)
 
-    adapted_samples_and_conditions = adapter(samples_and_conditions, strict=False)
+    # remove inference_variables from sample output and apply adapter
+    inference_variables_keys = approximator.sample(num_samples=num_samples, conditions=data).keys()
+    for key in inference_variables_keys:
+        samples_and_conditions.pop(key)
+    adapted_conditions = adapter(samples_and_conditions, strict=False)
 
-    assert any(k in adapted_samples_and_conditions for k in approximator.CONDITION_KEYS), (
+    assert any(k in adapted_conditions for k in approximator.CONDITION_KEYS), (
         f"adapter(approximator.sample(..., keep_conditions=True)) must return at least one of"
-        f"{approximator.CONDITION_KEYS!r}. Keys are {adapted_samples_and_conditions.keys()}."
+        f"{approximator.CONDITION_KEYS!r}. Keys are {adapted_conditions.keys()}."
     )
 
-    for key, value in adapted_samples_and_conditions.items():
+    for key, value in adapted_conditions.items():
         assert value.shape[:2] == (num_batches * batch_size, num_samples), (
             f"{key} should have shape ({num_batches * batch_size}, {num_samples}, ...) but has {value.shape}."
         )