Improve coupling flow interface and tests

Author: stefanradev93

bayesflow/networks/coupling_flow/actnorm.py

@@ -34,7 +34,7 @@ def build(self, xz_shape: Shape, **kwargs):
         self.scale = self.add_weight(shape=(xz_shape[-1],), initializer="ones", name="scale")
         self.bias = self.add_weight(shape=(xz_shape[-1],), initializer="zeros", name="bias")
 
-    def call(self, xz: Tensor, inverse: bool = False, **kwargs):
+    def call(self, xz: Tensor, inverse: bool = False, **kwargs) -> (Tensor, Tensor):
         if inverse:
             return self._inverse(xz, **kwargs)
         return self._forward(xz, **kwargs)

bayesflow/networks/coupling_flow/coupling_flow.py

@@ -46,6 +46,8 @@ def __init__(
         permutation: str | None = "random",
         use_actnorm: bool = True,
         base_distribution: str = "normal",
+        subnet_kwargs: dict[str, any] = None,
+        transform_kwargs: dict[str, any] = None,
         **kwargs,
     ):
         """
@@ -82,9 +84,15 @@ def __init__(
         base_distribution : str, optional
             The base probability distribution from which samples are drawn, such as
             "normal". Default is "normal".
+        subnet_kwargs : dict of str to any, optional
+            Keyword arguments forwarded to the subnet (e.g., MLP) constructor within
+            each coupling layer, such as hidden sizes or activation choices.
+        transform_kwargs : dict of str to any, optional
+            Keyword arguments forwarded to the affine or spline transforms
+            (e.g., bins for splines)
         **kwargs
-            Additional keyword arguments passed to the ActNorm, permutation, and
-            coupling layers for customization.
+            Additional keyword arguments passed to `InvertibleLayer`.
+
         """
         super().__init__(base_distribution=base_distribution, **kwargs)
 
@@ -97,12 +105,18 @@ def __init__(
         self.invertible_layers = []
         for i in range(depth):
             if use_actnorm:
-                self.invertible_layers.append(ActNorm(**kwargs.get("actnorm_kwargs", {})))
+                self.invertible_layers.append(ActNorm())
 
-            if (p := find_permutation(permutation, **kwargs.get("permutation_kwargs", {}))) is not None:
+            if (p := find_permutation(permutation)) is not None:
                 self.invertible_layers.append(p)
 
-            self.invertible_layers.append(DualCoupling(subnet, transform, **kwargs.get("coupling_kwargs", {})))
+            self.invertible_layers.append(
+                DualCoupling(subnet, transform, subnet_kwargs=subnet_kwargs, transform_kwargs=transform_kwargs)
+            )
+
+        # We only need to do this from coupling flows, since we do not serialize invertible layers
+        self.subnet_kwargs = subnet_kwargs
+        self.transform_kwargs = transform_kwargs
 
     # noinspection PyMethodOverriding
     def build(self, xz_shape, conditions_shape=None):
@@ -126,6 +140,8 @@ def get_config(self):
             "permutation": self.permutation,
             "use_actnorm": self.use_actnorm,
             "base_distribution": self.base_distribution,
+            "subnet_kwargs": self.subnet_kwargs,
+            "transform_kwargs": self.transform_kwargs,
         }
 
         return base_config | serialize(config)

bayesflow/networks/coupling_flow/couplings/single_coupling.py

@@ -29,16 +29,20 @@ def __init__(
         self,
         subnet: str | type = "mlp",
         transform: str = "affine",
+        subnet_kwargs: dict[str, any] = None,
+        transform_kwargs: dict[str, any] = None,
         **kwargs,
     ):
         super().__init__(**kwargs)
 
-        subnet_kwargs = kwargs.get("subnet_kwargs", {})
+        subnet_kwargs = subnet_kwargs or {}
+        transform_kwargs = transform_kwargs or {}
+
         if subnet == "mlp":
             subnet_kwargs = SingleCoupling.MLP_DEFAULT_CONFIG | subnet_kwargs
 
         self.subnet = find_network(subnet, **subnet_kwargs)
-        self.transform = find_transform(transform, **kwargs.get("transform_kwargs", {}))
+        self.transform = find_transform(transform, **transform_kwargs)
 
         self.output_projector = keras.layers.Dense(
             units=None, kernel_initializer="zeros", bias_initializer="zeros", name="output_projector"

bayesflow/networks/coupling_flow/permutations/fixed_permutation.py

@@ -1,16 +1,22 @@
 import keras
-from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Shape, Tensor
+from bayesflow.utils.serialization import serializable
+
 from ..invertible_layer import InvertibleLayer
 
 
-@serializable(package="networks.coupling_flow")
+@serializable
 class FixedPermutation(InvertibleLayer):
-    def __init__(self, forward_indices=None, inverse_indices=None, **kwargs):
+    """
+    Interface class for permutations with no learnable parameters. Child classes should
+    create forward and inverse indices in the associated build() method.
+    """
+
+    def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        self.forward_indices = forward_indices
-        self.inverse_indices = inverse_indices
+        self.forward_indices = None
+        self.inverse_indices = None
 
     def call(self, xz: Tensor, inverse: bool = False, **kwargs):
         if inverse:

bayesflow/networks/coupling_flow/permutations/orthogonal.py

@@ -1,11 +1,12 @@
 from keras import ops
-from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Shape, Tensor
+from bayesflow.utils.serialization import serializable
+
 from ..invertible_layer import InvertibleLayer
 
 
-@serializable(package="networks.coupling_flow")
+@serializable
 class OrthogonalPermutation(InvertibleLayer):
     """Implements a learnable orthogonal transformation according to [1]. Can be
     used as an alternative to a fixed ``Permutation`` layer.
@@ -21,7 +22,7 @@ def __init__(self, **kwargs):
     def build(self, xz_shape: Shape, **kwargs) -> None:
         self.weight = self.add_weight(shape=(xz_shape[-1], xz_shape[-1]), initializer="orthogonal", trainable=True)
 
-    def call(self, xz: Tensor, inverse: bool = False, **kwargs):
+    def call(self, xz: Tensor, inverse: bool = False, **kwargs) -> (Tensor, Tensor):
         if inverse:
             return self._inverse(xz)
         return self._forward(xz)

bayesflow/networks/coupling_flow/permutations/random.py

@@ -1,11 +1,12 @@
 import keras
-from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Shape
+from bayesflow.utils.serialization import serializable
+
 from .fixed_permutation import FixedPermutation
 
 
-@serializable(package="networks.coupling_flow")
+@serializable
 class RandomPermutation(FixedPermutation):
     # noinspection PyMethodOverriding
     def build(self, xz_shape: Shape, **kwargs) -> None:

bayesflow/networks/coupling_flow/permutations/swap.py

@@ -1,11 +1,12 @@
 import keras
-from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Shape
+from bayesflow.utils.serialization import serializable
+
 from .fixed_permutation import FixedPermutation
 
 
-@serializable(package="networks.coupling_flow")
+@serializable
 class Swap(FixedPermutation):
     def build(self, xz_shape: Shape, **kwargs) -> None:
         shift = xz_shape[-1] // 2

bayesflow/networks/coupling_flow/transforms/affine_transform.py

@@ -1,12 +1,13 @@
 import keras.ops as ops
-from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Tensor
 from bayesflow.utils.keras_utils import shifted_softplus
+from bayesflow.utils.serialization import serializable
+
 from .transform import Transform
 
 
-@serializable(package="networks.coupling_flow")
+@serializable
 class AffineTransform(Transform):
     def __init__(self, clamp: bool = True, **kwargs):
         super().__init__(**kwargs)

bayesflow/networks/coupling_flow/transforms/spline_transform.py

@@ -1,17 +1,16 @@
 import keras
 import numpy as np
-from keras.saving import (
-    register_keras_serializable as serializable,
-)
 
 from bayesflow.types import Tensor
 from bayesflow.utils import pad, searchsorted
 from bayesflow.utils.keras_utils import shifted_softplus
+from bayesflow.utils.serialization import serializable
+
 from ._rational_quadratic import _rational_quadratic_spline
 from .transform import Transform
 
 
-@serializable(package="networks.coupling_flow")
+@serializable
 class SplineTransform(Transform):
     def __init__(
         self,

tests/test_networks/conftest.py

@@ -8,23 +8,41 @@ def flow_matching():
     from bayesflow.networks import FlowMatching
 
     return FlowMatching(
-        subnet=MLP([64, 64]),
+        subnet=MLP([8, 8]),
         integrate_kwargs={"method": "rk45", "steps": 100},
     )
 
 
 @pytest.fixture()
-def coupling_flow():
+def consistency_model():
+    from bayesflow.networks import ConsistencyModel
+
+    return ConsistencyModel(total_steps=100, subnet=MLP([8, 8]))
+
+
+@pytest.fixture()
+def affine_coupling_flow():
+    from bayesflow.networks import CouplingFlow
+
+    return CouplingFlow(
+        depth=2, subnet="mlp", subnet_kwargs=dict(widths=[8, 8]), transform="affine", transform_kwargs=dict(clamp=1.8)
+    )
+
+
+@pytest.fixture()
+def spline_coupling_flow():
     from bayesflow.networks import CouplingFlow
 
-    return CouplingFlow(depth=2, subnet="mlp", subnet_kwargs=dict(widths=[64, 64]))
+    return CouplingFlow(
+        depth=2, subnet="mlp", subnet_kwargs=dict(widths=[8, 8]), transform="spline", transform_kwargs=dict(bins=8)
+    )
 
 
 @pytest.fixture()
 def free_form_flow():
     from bayesflow.experimental import FreeFormFlow
 
-    return FreeFormFlow(encoder_subnet=MLP([64, 64]), decoder_subnet=MLP([64, 64]))
+    return FreeFormFlow(encoder_subnet=MLP([16, 16]), decoder_subnet=MLP([16, 16]))
 
 
 @pytest.fixture()
@@ -47,7 +65,7 @@ def typical_point_inference_network_subnet():
     from bayesflow.networks import PointInferenceNetwork
     from bayesflow.scores import MeanScore, MedianScore, QuantileScore, MultivariateNormalScore
 
-    subnet = MLP([64, 64])
+    subnet = MLP([16, 8])
 
     return PointInferenceNetwork(
         scores=dict(
@@ -61,7 +79,14 @@ def typical_point_inference_network_subnet():
 
 
 @pytest.fixture(
-    params=["typical_point_inference_network", "coupling_flow", "flow_matching", "free_form_flow"], scope="function"
+    params=[
+        "typical_point_inference_network",
+        "affine_coupling_flow",
+        "spline_coupling_flow",
+        "flow_matching",
+        "free_form_flow",
+    ],
+    scope="function",
 )
 def inference_network(request):
     return request.getfixturevalue(request.param)
@@ -80,7 +105,9 @@ def inference_network_subnet(request):
     return request.getfixturevalue(request.param)
 
 
-@pytest.fixture(params=["coupling_flow", "flow_matching", "free_form_flow"], scope="function")
+@pytest.fixture(
+    params=["affine_coupling_flow", "spline_coupling_flow", "flow_matching", "free_form_flow"], scope="function"
+)
 def generative_inference_network(request):
     return request.getfixturevalue(request.param)