change to simpler separate test

The cost of the continuous models on the CI is too high

Author: vpratz

tests/test_networks/conftest.py

@@ -61,12 +61,12 @@ def get_config(self):
 
 
 @pytest.fixture()
-def diffusion_model_edm_F_subnet_concatenate():
+def diffusion_model_edm_F_subnet_separate_inputs():
     from bayesflow.networks import DiffusionModel
 
     return DiffusionModel(
         subnet=ConcatenateMLP([8, 8]),
-        integrate_kwargs={"method": "rk45", "steps": 250},
+        integrate_kwargs={"method": "rk45", "steps": 4},
         noise_schedule="edm",
         prediction_type="F",
         concatenate_subnet_input=False,
@@ -144,11 +144,11 @@ def flow_matching():
 
 
 @pytest.fixture()
-def flow_matching_subnet_concatenate():
+def flow_matching_subnet_separate_inputs():
     from bayesflow.networks import FlowMatching
 
     return FlowMatching(
-        subnet=ConcatenateMLP([8, 8]), integrate_kwargs={"method": "rk45", "steps": 100}, concatenate_subnet_input=False
+        subnet=ConcatenateMLP([8, 8]), integrate_kwargs={"method": "rk45", "steps": 4}, concatenate_subnet_input=False
     )
 
 
@@ -160,10 +160,10 @@ def consistency_model():
 
 
 @pytest.fixture()
-def consistency_model_subnet_concatenate():
+def consistency_model_subnet_separate_inputs():
     from bayesflow.networks import ConsistencyModel
 
-    return ConsistencyModel(total_steps=100, subnet=ConcatenateMLP([8, 8]), concatenate_subnet_input=False)
+    return ConsistencyModel(total_steps=4, subnet=ConcatenateMLP([8, 8]), concatenate_subnet_input=False)
 
 
 @pytest.fixture()
@@ -230,12 +230,9 @@ def typical_point_inference_network_subnet():
         "affine_coupling_flow",
         "spline_coupling_flow",
         "flow_matching",
-        pytest.param("flow_matching_subnet_concatenate"),
         "free_form_flow",
         "consistency_model",
-        pytest.param("consistency_model_subnet_concatenate"),
         pytest.param("diffusion_model_edm_F"),
-        pytest.param("diffusion_model_edm_F_subnet_concatenate"),
         pytest.param("diffusion_model_edm_noise", marks=pytest.mark.slow),
         pytest.param("diffusion_model_cosine_velocity", marks=pytest.mark.slow),
         pytest.param("diffusion_model_cosine_F", marks=pytest.mark.slow),
@@ -263,16 +260,12 @@ def inference_network_subnet(request):
 
 @pytest.fixture(
     params=[
-        pytest.param("diffusion_model_edm_F_subnet_concatenate"),
         "affine_coupling_flow",
         "spline_coupling_flow",
         "flow_matching",
-        pytest.param("flow_matching_subnet_concatenate"),
         "free_form_flow",
         "consistency_model",
-        pytest.param("consistency_model_subnet_concatenate"),
         pytest.param("diffusion_model_edm_F"),
-        pytest.param("diffusion_model_edm_F_subnet_concatenate"),
         pytest.param(
             "diffusion_model_edm_noise",
             marks=[
@@ -309,6 +302,18 @@ def generative_inference_network(request):
     return request.getfixturevalue(request.param)
 
 
+@pytest.fixture(
+    params=[
+        pytest.param("flow_matching_subnet_separate_inputs"),
+        pytest.param("consistency_model_subnet_separate_inputs"),
+        pytest.param("diffusion_model_edm_F_subnet_separate_inputs"),
+    ],
+    scope="function",
+)
+def inference_network_subnet_separate_inputs(request):
+    return request.getfixturevalue(request.param)
+
+
 @pytest.fixture(scope="function")
 def time_series_network(summary_dim):
     from bayesflow.networks import TimeSeriesNetwork

tests/test_networks/test_inference_networks.py

@@ -162,3 +162,16 @@ def test_compute_metrics(inference_network, random_samples, random_conditions):
 
     metrics = inference_network.compute_metrics(random_samples, conditions=random_conditions)
     assert "loss" in metrics
+
+
+def test_subnet_separate_inputs(inference_network_subnet_separate_inputs, random_samples, random_conditions):
+    xz_shape = keras.ops.shape(random_samples)
+    conditions_shape = keras.ops.shape(random_conditions) if random_conditions is not None else None
+    inference_network_subnet_separate_inputs.build(xz_shape, conditions_shape)
+
+    assert inference_network_subnet_separate_inputs.built is True
+
+    # check the model has variables
+    assert inference_network_subnet_separate_inputs.variables, "Model has no variables."
+
+    inference_network_subnet_separate_inputs(random_samples, random_conditions, inverse=True)