Simplify import structure

bayesflow-org · Jun 12, 2024 · 060fd10 · 060fd10
1 parent 7280f40
commit 060fd10
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 3 deletions.
diff --git a/bayesflow/experimental/networks/coupling_flow/coupling_flow.py b/bayesflow/experimental/networks/coupling_flow/coupling_flow.py
@@ -6,6 +6,7 @@
 
 from bayesflow.experimental.types import Tensor
 from bayesflow.experimental.utils import find_permutation, keras_kwargs
+
 from .actnorm import ActNorm
 from .couplings import DualCoupling
 from ..inference_network import InferenceNetwork

diff --git a/bayesflow/experimental/networks/coupling_flow/transforms/affine_transform.py b/bayesflow/experimental/networks/coupling_flow/transforms/affine_transform.py
@@ -1,7 +1,6 @@
+
 import keras.ops as ops
-from keras.saving import (
-    register_keras_serializable,
-)
+from keras.saving import register_keras_serializable
 
 from bayesflow.experimental.types import Tensor
 from .transform import Transform

diff --git a/bayesflow/experimental/networks/coupling_flow/transforms/spline_transform.py b/bayesflow/experimental/networks/coupling_flow/transforms/spline_transform.py
@@ -2,11 +2,13 @@
 import math
 
 from keras import ops
+from keras.saving import register_keras_serializable
 
 from bayesflow.experimental.types import Tensor
 from .transform import Transform
 
 
+@register_keras_serializable(package="bayesflow.networks.coupling_flow")
 class SplineTransform(Transform):
     def __init__(self, bins=16, default_domain=(-5., 5., -5., 5.), **kwargs):
         super().__init__(**kwargs)

diff --git a/bayesflow/experimental/networks/coupling_flow/transforms/transform.py b/bayesflow/experimental/networks/coupling_flow/transforms/transform.py
@@ -1,3 +1,4 @@
+
 from bayesflow.experimental.types import Tensor
 from ..invertible_layer import InvertibleLayer