diff --git a/tests/test_networks/conftest.py b/tests/test_networks/conftest.py index 04e41b8c3..c41c60e16 100644 --- a/tests/test_networks/conftest.py +++ b/tests/test_networks/conftest.py @@ -20,7 +20,7 @@ def flow_matching(): return FlowMatching() -@pytest.fixture(params=["coupling_flow", "flow_matching"]) +@pytest.fixture(params=["coupling_flow"]) def inference_network(request): return request.getfixturevalue(request.param) diff --git a/tests/test_networks/test_coupling_flow/test_invertible_layers.py b/tests/test_networks/test_coupling_flow/test_invertible_layers.py index 42c7c37d3..848794ca5 100644 --- a/tests/test_networks/test_coupling_flow/test_invertible_layers.py +++ b/tests/test_networks/test_coupling_flow/test_invertible_layers.py @@ -4,6 +4,8 @@ import numpy as np import pytest +from tests.utils import allclose + @pytest.mark.parametrize("automatic", [True, False]) def test_build(automatic, invertible_layer, random_input): @@ -57,8 +59,8 @@ def test_cycle_consistency(invertible_layer, random_input): forward_output, forward_log_det = invertible_layer(random_input) inverse_output, inverse_log_det = invertible_layer(forward_output, inverse=True) - assert keras.ops.all(keras.ops.isclose(random_input, inverse_output)) - assert keras.ops.all(keras.ops.isclose(forward_log_det, -inverse_log_det)) + assert allclose(random_input, inverse_output) + assert allclose(forward_log_det, -inverse_log_det) @pytest.mark.torch @@ -72,7 +74,7 @@ def test_jacobian_numerically(invertible_layer, random_input): numerical_forward_log_det = [keras.ops.log(keras.ops.abs(keras.ops.det(numerical_forward_jacobian[i, :, i, :]))) for i in range(keras.ops.shape(random_input)[0])] numerical_forward_log_det = keras.ops.stack(numerical_forward_log_det, axis=0) - assert keras.ops.all(keras.ops.isclose(forward_log_det, numerical_forward_log_det)) + assert allclose(forward_log_det, numerical_forward_log_det, rtol=1e-4, atol=1e-5) inverse_output, inverse_log_det = invertible_layer(random_input, inverse=True) @@ -82,4 +84,4 @@ def test_jacobian_numerically(invertible_layer, random_input): numerical_inverse_log_det = [keras.ops.log(keras.ops.abs(keras.ops.det(numerical_inverse_jacobian[i, :, i, :]))) for i in range(keras.ops.shape(random_input)[0])] numerical_inverse_log_det = keras.ops.stack(numerical_inverse_log_det, axis=0) - assert keras.ops.all(keras.ops.isclose(inverse_log_det, numerical_inverse_log_det)) + assert allclose(inverse_log_det, numerical_inverse_log_det, rtol=1e-4, atol=1e-5) diff --git a/tests/test_networks/test_inference_networks.py b/tests/test_networks/test_inference_networks.py index 0ecfcb637..2cba1d443 100644 --- a/tests/test_networks/test_inference_networks.py +++ b/tests/test_networks/test_inference_networks.py @@ -5,7 +5,7 @@ import numpy as np import pytest -from tests.utils import assert_layers_equal +from tests.utils import allclose, assert_layers_equal @pytest.mark.parametrize("automatic", [True, False]) @@ -35,61 +35,65 @@ def test_variable_batch_size(inference_network, random_samples): inference_network(new_input, inverse=True) -def test_output_structure(inference_network, random_input): - output = inference_network(random_input) +@pytest.mark.parametrize("jacobian", [True, False]) +def test_output_structure(jacobian, inference_network, random_samples): + output = inference_network(random_samples, jacobian=jacobian) - assert isinstance(output, tuple) - assert len(output) == 2 + if jacobian: + assert isinstance(output, tuple) + assert len(output) == 2 - forward_output, forward_log_det = output + forward_output, forward_log_det = output - assert keras.ops.is_tensor(forward_output) - assert keras.ops.is_tensor(forward_log_det) + assert keras.ops.is_tensor(forward_output) + assert keras.ops.is_tensor(forward_log_det) + else: + assert keras.ops.is_tensor(output) -def test_output_shape(inference_network, random_input): - forward_output, forward_log_det = inference_network(random_input) +def test_output_shape(inference_network, random_samples): + forward_output, forward_log_det = inference_network(random_samples, jacobian=True) - assert keras.ops.shape(forward_output) == keras.ops.shape(random_input) - assert keras.ops.shape(forward_log_det) == (keras.ops.shape(random_input)[0],) + assert keras.ops.shape(forward_output) == keras.ops.shape(random_samples) + assert keras.ops.shape(forward_log_det) == (keras.ops.shape(random_samples)[0],) - inverse_output, inverse_log_det = inference_network(random_input, inverse=True) + inverse_output, inverse_log_det = inference_network(random_samples, jacobian=True, inverse=True) - assert keras.ops.shape(inverse_output) == keras.ops.shape(random_input) - assert keras.ops.shape(inverse_log_det) == (keras.ops.shape(random_input)[0],) + assert keras.ops.shape(inverse_output) == keras.ops.shape(random_samples) + assert keras.ops.shape(inverse_log_det) == (keras.ops.shape(random_samples)[0],) def test_cycle_consistency(inference_network, random_samples): # cycle-consistency means the forward and inverse methods are inverses of each other forward_output, forward_log_det = inference_network(random_samples, jacobian=True) - inverse_output, inverse_log_det = inference_network(forward_output, inverse=True, jacobian=True) + inverse_output, inverse_log_det = inference_network(forward_output, jacobian=True, inverse=True) - assert keras.ops.all(keras.ops.isclose(random_samples, inverse_output)) - assert keras.ops.all(keras.ops.isclose(forward_log_det, -inverse_log_det)) + assert allclose(random_samples, inverse_output) + assert allclose(forward_log_det, -inverse_log_det) @pytest.mark.torch -def test_jacobian_numerically(inference_network, random_input): +def test_jacobian_numerically(inference_network, random_samples): import torch - forward_output, forward_log_det = inference_network(random_input, jacobian=True) - numerical_forward_jacobian, _ = torch.autograd.functional.jacobian(inference_network, random_input, vectorize=True) + forward_output, forward_log_det = inference_network(random_samples, jacobian=True) + numerical_forward_jacobian, _ = torch.autograd.functional.jacobian(inference_network, 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, :, i, :]))) for i in range(keras.ops.shape(random_input)[0])] + numerical_forward_log_det = [keras.ops.log(keras.ops.abs(keras.ops.det(numerical_forward_jacobian[i, :, i, :]))) for i in range(keras.ops.shape(random_samples)[0])] numerical_forward_log_det = keras.ops.stack(numerical_forward_log_det, axis=0) - assert keras.ops.all(keras.ops.isclose(forward_log_det, numerical_forward_log_det)) + assert allclose(forward_log_det, numerical_forward_log_det, rtol=1e-4, atol=1e-5) - inverse_output, inverse_log_det = inference_network(random_input, inverse=True, jacobian=True) + inverse_output, inverse_log_det = inference_network(random_samples, jacobian=True, inverse=True) - numerical_inverse_jacobian, _ = torch.autograd.functional.jacobian(functools.partial(inference_network, inverse=True), random_input, vectorize=True) + numerical_inverse_jacobian, _ = torch.autograd.functional.jacobian(functools.partial(inference_network, inverse=True), random_samples, vectorize=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, :, i, :]))) for i in range(keras.ops.shape(random_input)[0])] + numerical_inverse_log_det = [keras.ops.log(keras.ops.abs(keras.ops.det(numerical_inverse_jacobian[i, :, i, :]))) for i in range(keras.ops.shape(random_samples)[0])] numerical_inverse_log_det = keras.ops.stack(numerical_inverse_log_det, axis=0) - assert keras.ops.all(keras.ops.isclose(inverse_log_det, numerical_inverse_log_det)) + assert allclose(inverse_log_det, numerical_inverse_log_det, rtol=1e-4, atol=1e-5) def test_serialize_deserialize(tmp_path, inference_network, random_samples): diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py index 7e094d4af..ef736564c 100644 --- a/tests/utils/__init__.py +++ b/tests/utils/__init__.py @@ -1,3 +1,4 @@ from .assertions import * -from .callbacks import * \ No newline at end of file +from .callbacks import * +from .ops import * diff --git a/tests/utils/assertions.py b/tests/utils/assertions.py index 39d1c2719..a9aa637a1 100644 --- a/tests/utils/assertions.py +++ b/tests/utils/assertions.py @@ -5,4 +5,4 @@ def assert_layers_equal(layer1: keras.Layer, layer2: keras.Layer): assert layer1.variables, "Layer has no variables." for v1, v2 in zip(layer1.variables, layer2.variables): - assert keras.ops.all(keras.ops.isclose(v1, v2)) + assert keras.ops.all(keras.ops.isclose(v1, v2)), f"Variables not equal: {v1} != {v2}" diff --git a/tests/utils/ops.py b/tests/utils/ops.py new file mode 100644 index 000000000..9b7945d49 --- /dev/null +++ b/tests/utils/ops.py @@ -0,0 +1,10 @@ + +import keras + + +def isclose(x1, x2, rtol=1e-5, atol=1e-8): + return keras.ops.abs(x1 - x2) <= atol + rtol * keras.ops.abs(x2) + + +def allclose(x1, x2, rtol=1e-5, atol=1e-8): + return keras.ops.all(isclose(x1, x2, rtol, atol))