test improvements

Author: LarsKue

tests/test_networks/test_inference_networks.py

@@ -6,7 +6,7 @@
     serialize_keras_object as serialize,
 )
 
-from tests.utils import allclose, assert_layers_equal
+from tests.utils import assert_allclose, assert_layers_equal
 
 
 def test_build(inference_network, random_samples, random_conditions):
@@ -83,58 +83,30 @@ def test_cycle_consistency(generative_inference_network, random_samples, random_
         forward_output, conditions=random_conditions, density=True, inverse=True
     )
 
-    assert allclose(random_samples, inverse_output, atol=1e-3, rtol=1e-3)
-    assert allclose(forward_log_density, inverse_log_density, atol=1e-3, rtol=1e-3)
+    assert_allclose(random_samples, inverse_output, atol=1e-3, rtol=1e-3)
+    assert_allclose(forward_log_density, inverse_log_density, atol=1e-3, rtol=1e-3)
 
 
-# TODO: make this backend-agnostic
-@pytest.mark.torch
 def test_density_numerically(generative_inference_network, random_samples, random_conditions):
-    import torch
+    from bayesflow.utils import jacobian
 
-    forward_output, forward_log_density = generative_inference_network(
-        random_samples, conditions=random_conditions, density=True
-    )
+    output, log_density = generative_inference_network(random_samples, conditions=random_conditions, density=True)
 
     def f(x):
         return generative_inference_network(x, conditions=random_conditions)
 
-    numerical_forward_jacobian, *_ = torch.autograd.functional.jacobian(f, random_samples, vectorize=True)
-
-    # TODO: torch is somehow permuted wrt keras
-    numerical_forward_log_det = [
-        keras.ops.log(keras.ops.abs(keras.ops.det(numerical_forward_jacobian[:, i, :])))
-        for i in range(keras.ops.shape(random_samples)[0])
-    ]
-    numerical_forward_log_det = keras.ops.stack(numerical_forward_log_det, axis=0)
-
-    log_prob = generative_inference_network.base_distribution.log_prob(forward_output)
-
-    numerical_forward_log_density = log_prob + numerical_forward_log_det
-
-    assert allclose(forward_log_density, numerical_forward_log_density, rtol=1e-4, atol=1e-5)
-
-    inverse_output, inverse_log_density = generative_inference_network(
-        random_samples, conditions=random_conditions, density=True, inverse=True
-    )
-
-    def f(x):
-        return generative_inference_network(x, conditions=random_conditions, inverse=True)
-
-    numerical_inverse_jacobian, *_ = torch.autograd.functional.jacobian(f, random_samples, vectorize=True)
+    numerical_output, numerical_jacobian = jacobian(f, random_samples, return_output=True)
 
-    # TODO: torch is somehow permuted wrt keras
-    numerical_inverse_log_det = [
-        keras.ops.log(keras.ops.abs(keras.ops.det(numerical_inverse_jacobian[:, i, :])))
-        for i in range(keras.ops.shape(random_samples)[0])
-    ]
-    numerical_inverse_log_det = keras.ops.stack(numerical_inverse_log_det, axis=0)
+    # output should be identical, otherwise this test does not work (e.g. for stochastic networks)
+    assert keras.ops.all(keras.ops.isclose(output, numerical_output))
 
-    log_prob = generative_inference_network.base_distribution.log_prob(random_samples)
+    log_prob = generative_inference_network.base_distribution.log_prob(output)
 
-    numerical_inverse_log_density = log_prob - numerical_inverse_log_det
+    # use change of variables to compute the numerical log density
+    numerical_log_density = log_prob + keras.ops.log(keras.ops.abs(keras.ops.det(numerical_jacobian)))
 
-    assert allclose(inverse_log_density, numerical_inverse_log_density, rtol=1e-4, atol=1e-5)
+    # use a high tolerance because the numerical jacobian is not very accurate
+    assert_allclose(log_density, numerical_log_density, rtol=1e-3, atol=1e-3)
 
 
 def test_serialize_deserialize(inference_network_subnet, subnet, random_samples, random_conditions):