Merge branch 'streamlined-backend' of https://github.com/stefanradev9…

…3/BayesFlow into streamlined-backend

bayesflow/experimental/networks/coupling_flow/couplings/single_coupling.py

@@ -0,0 +1,60 @@
+
+import keras
+from keras.saving import (
+    register_keras_serializable
+)
+
+from bayesflow.experimental.types import Tensor
+from bayesflow.experimental.utils import find_network
+from ..invertible_layer import InvertibleLayer
+from ..transforms import find_transform
+
+
+@register_keras_serializable(package="bayesflow.networks.coupling_flow")
+class SingleCoupling(InvertibleLayer):
+    """
+    Implements a single coupling layer as a composition of a subnet and a transform.
+
+    Subnet output tensors are linearly mapped to the correct dimension.
+    """
+    def __init__(self, network: str = "resnet", transform: str = "affine", **kwargs):
+        super().__init__(**kwargs)
+        self.dense = keras.layers.Dense(None, kernel_initializer="zeros", bias_initializer="zeros")
+        self.network = find_network(network)
+        self.transform = find_transform(transform)
+
+    # noinspection PyMethodOverriding
+    def build(self, x1_shape, x2_shape):
+        self.dense.units = self.transform.params_per_dim * x2_shape[-1]
+
+    def call(self, x1: Tensor, x2: Tensor, conditions: any = None, inverse: bool = False) -> ((Tensor, Tensor), Tensor):
+        if inverse:
+            return self._inverse(x1, x2, conditions=conditions)
+        return self._forward(x1, x2, conditions=conditions)
+
+    def _forward(self, x1: Tensor, x2: Tensor, conditions: any = None) -> ((Tensor, Tensor), Tensor):
+        """ Transform (x1, x2) -> (x1, f(x2; x1)) """
+        z1 = x1
+        parameters = self.get_parameters(x1, conditions)
+        z2, log_det = self.transform(x2, parameters=parameters)
+
+        return (z1, z2), log_det
+
+    def _inverse(self, z1: Tensor, z2: Tensor, conditions: any = None) -> ((Tensor, Tensor), Tensor):
+        """ Transform (x1, f(x2; x1)) -> (x1, x2) """
+        x1 = z1
+        parameters = self.get_parameters(x1, conditions)
+        x2, log_det = self.transform(z2, parameters=parameters, inverse=True)
+
+        return (x1, x2), log_det
+
+    def get_parameters(self, x, conditions: any = None) -> dict[str, Tensor]:
+        # TODO: pass conditions to subnet via kwarg if possible
+        if keras.ops.is_tensor(conditions):
+            x = keras.ops.concatenate([x, conditions], axis=-1)
+
+        parameters = self.dense(self.network(x))
+        parameters = self.transform.split_parameters(parameters)
+        parameters = self.transform.constrain_parameters(parameters)
+
+        return parameters

bayesflow/experimental/networks/coupling_flow/permutations/__init__.py

@@ -1,3 +1,4 @@
 
 from .orthogonal import OrthogonalPermutation
+from .fixed_permutation import FixedPermutation
 from .swap import Swap

bayesflow/experimental/networks/coupling_flow/permutations/fixed_permutation.py

@@ -15,7 +15,7 @@ def __init__(self, forward_indices=None, inverse_indices=None, **kwargs):
         self.forward_indices = forward_indices
         self.inverse_indices = inverse_indices
 
-    def call(self, xz: Tensor, inverse: bool = False):
+    def call(self, xz: Tensor, inverse: bool = False, **kwargs):
         if inverse:
             return self._inverse(xz)
         return self._forward(xz)
@@ -25,10 +25,10 @@ def build(self, input_shape: Shape) -> None:
 
     def _forward(self, x: Tensor) -> (Tensor, Tensor):
         z = keras.ops.take(x, self.forward_indices, axis=-1)
-        log_det = keras.ops.zeros(keras.ops.shape(x)[0])
+        log_det = 0.
         return z, log_det
 
     def _inverse(self, z: Tensor) -> (Tensor, Tensor):
         x = keras.ops.take(z, self.inverse_indices, axis=-1)
-        log_det = keras.ops.zeros(keras.ops.shape(z)[0])
+        log_det = 0.
         return x, log_det

bayesflow/experimental/networks/coupling_flow/permutations/orthogonal.py

@@ -1,10 +1,14 @@
 
 from keras import ops
+from keras.saving import (
+    register_keras_serializable,
+)
 
 from bayesflow.experimental.types import Shape, Tensor
 from ..invertible_layer import InvertibleLayer
 
 
+@register_keras_serializable(package="bayesflow.networks.coupling_flow")
 class OrthogonalPermutation(InvertibleLayer):
     """Implements a learnable orthogonal transformation according to [1]. Can be
     used as an alternative to a fixed ``Permutation`` layer.
@@ -25,7 +29,7 @@ def build(self, input_shape: Shape) -> None:
             trainable=True
         )
 
-    def call(self, xz: Tensor, inverse: bool = False):
+    def call(self, xz: Tensor, inverse: bool = False, **kwargs):
         if inverse:
             return self._inverse(xz)
         return self._forward(xz)

bayesflow/experimental/networks/coupling_flow/permutations/random.py

@@ -14,13 +14,13 @@ def build(self, input_shape: Shape) -> None:
         forward_indices = keras.random.shuffle(keras.ops.arange(input_shape[-1]))
         inverse_indices = keras.ops.argsort(forward_indices)
 
-        self.forward_indices = self.add_variable(
+        self.forward_indices = self.add_weight(
             shape=(input_shape[-1],),
             initializer=keras.initializers.Constant(forward_indices),
             trainable=False
         )
 
-        self.inverse_indices = self.add_variable(
+        self.inverse_indices = self.add_weight(
             shape=(input_shape[-1],),
             initializer=keras.initializers.Constant(inverse_indices),
             trainable=False

bayesflow/experimental/networks/coupling_flow/permutations/swap.py

@@ -1,5 +1,8 @@
 
 import keras
+from keras.saving import (
+    register_keras_serializable,
+)
 
 from bayesflow.experimental.types import Shape
 from .fixed_permutation import FixedPermutation