adapt tests for torch

Author: vpratz

tests/test_compatibility/test_approximators/test_point_approximator/conftest.py

@@ -6,11 +6,6 @@ def batch_size():
     return 8
 
 
-@pytest.fixture()
-def num_samples():
-    return 100
-
-
 @pytest.fixture(params=["single_parametric", "multiple_parametric"])
 def point_inference_network(request):
     match request.param:
@@ -52,3 +47,13 @@ def approximator(adapter, point_inference_network, summary_network, standardize)
         summary_network=summary_network,
         standardize=standardize,
     )
+
+
+@pytest.fixture()
+def adapter():
+    from bayesflow import ContinuousApproximator
+
+    return ContinuousApproximator.build_adapter(
+        inference_variables=["mean", "std"],
+        inference_conditions=["x"],
+    )

tests/test_compatibility/test_approximators/test_point_approximator/test_point_approximator.py

@@ -5,10 +5,9 @@
 
 
 @pytest.mark.parametrize(
-    "summary_network,simulator,adapter,standardize",
+    "summary_network,simulator,standardize",
     [
-        ["deep_set", "sir", "summary", "all"],  # use deep_set for speed
-        [None, "two_moons", "direct", None],
+        [None, "normal", "all"],
     ],
     indirect=True,
 )

tests/test_compatibility/test_distributions/test_distributions.py

@@ -48,12 +48,20 @@ def from_config(cls, config, custom_objects=None):
 
             model = DummyModel(distribution)
             model.compile(loss=keras.losses.MeanSquaredError())
-            model.fit(
-                random_samples,
-                keras.ops.ones(keras.ops.shape(random_samples)[:-1]),
+            fit_kwargs = dict(
+                x=random_samples,
+                y=keras.ops.ones(keras.ops.shape(random_samples)[:-1]),
                 batch_size=keras.ops.shape(random_samples)[0],
                 epochs=1,
             )
+            if keras.backend.backend() == "torch":
+                import torch
+
+                with torch.enable_grad():
+                    model.fit(**fit_kwargs)
+            else:
+                model.fit(**fit_kwargs)
+
             model.save(filepaths["model"])
 
             output = self.evaluate(model.distribution, random_samples)

tests/test_compatibility/test_metrics/conftest.py

@@ -1,5 +1,5 @@
 import pytest
-import numpy as np
+import keras
 
 
 @pytest.fixture()
@@ -20,11 +20,9 @@ def metric(request):
 
 @pytest.fixture
 def samples_1():
-    rng = np.random.default_rng(seed=1)
-    return rng.normal(size=(2, 3)).astype(np.float32)
+    return keras.random.normal((2, 3), seed=1)
 
 
 @pytest.fixture
 def samples_2():
-    rng = np.random.default_rng(seed=2)
-    return rng.normal(size=(2, 3)).astype(np.float32)
+    return keras.random.normal((2, 3), seed=2)