Add FFJORD test

davidnabergoj · Sep 11, 2023 · 3c94a9b · 3c94a9b
1 parent 92a44d6
commit 3c94a9b
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/normalizing_flows/bijections/__init__.py b/normalizing_flows/bijections/__init__.py
@@ -10,4 +10,4 @@
     HouseholderSylvester, Sylvester
 from normalizing_flows.bijections.finite.residual.iterative import InvertibleResNet, ResFlow, \
     QuasiAutoregressiveFlow, ProximalResidualFlow
-from normalizing_flows.bijections.finite.linear import LowerTriangular, Orthogonal, LU, QR
+from normalizing_flows.bijections.finite.linear import LowerTriangular, Orthogonal, LU, QR
diff --git a/test/test_reconstruction_bijections.py b/test/test_reconstruction_bijections.py
@@ -8,6 +8,7 @@
 from normalizing_flows.bijections import RealNVP, MAF, CouplingRQNSF, MaskedAutoregressiveRQNSF, ResFlow, \
     InvertibleResNet, \
     ElementwiseAffine, ElementwiseShift, InverseAutoregressiveRQNSF, IAF, NICE
+from normalizing_flows.bijections import FFJORD
 from normalizing_flows.bijections.finite.base import ConditionalBijection
 from normalizing_flows.bijections.finite.residual.planar import Planar
 from normalizing_flows.bijections.finite.residual.radial import Radial
@@ -120,3 +121,15 @@ def test_masked_autoregressive(bijection_class: ConditionalBijection, batch_shap
 def test_residual(bijection_class: ConditionalBijection, batch_shape: Tuple, event_shape: Tuple, context_shape: Tuple):
     bijection, x, context = setup_data(bijection_class, batch_shape, event_shape, context_shape)
     assert_valid_reconstruction(bijection, x, context)
+
+
+@pytest.mark.parametrize('bijection_class', [
+    FFJORD
+])
+@pytest.mark.parametrize('batch_shape', __test_constants['batch_shape'])
+@pytest.mark.parametrize('event_shape', __test_constants['event_shape'])
+@pytest.mark.parametrize('context_shape', __test_constants['context_shape'])
+def test_continuous(bijection_class: ConditionalBijection, batch_shape: Tuple, event_shape: Tuple,
+                    context_shape: Tuple):
+    bijection, x, context = setup_data(bijection_class, batch_shape, event_shape, context_shape)
+    assert_valid_reconstruction(bijection, x, context)