serialization: apply new scheme for package (breaking change)

README.md

@@ -15,6 +15,12 @@ It provides users and researchers with:
 BayesFlow (version 2+) is designed to be a flexible and efficient tool that enables rapid statistical inference
 fueled by continuous progress in generative AI and Bayesian inference.
 
+> [!IMPORTANT]
+> As the 2.0 version introduced many new features, we still have to make breaking changes from time to time.
+> This especially concerns **saving and loading** of models. We aim to stabilize this from the 2.1 release onwards.
+> Until then, consider pinning your BayesFlow 2.0 installation to an exact version, or re-training after an update
+> for less costly models.
+
 ## Important Note for Existing Users
 
 You are currently looking at BayesFlow 2.0+, which is a complete rewrite of the library.

bayesflow/adapters/adapter.py

@@ -29,7 +29,7 @@
 from .transforms.filter_transform import Predicate
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Adapter(MutableSequence[Transform]):
     """
     Defines an adapter to apply various transforms to data.

bayesflow/adapters/transforms/as_set.py

@@ -5,7 +5,7 @@
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class AsSet(ElementwiseTransform):
     """The `.as_set(["x", "y"])` transform indicates that both `x` and `y` are treated as sets.
 

bayesflow/adapters/transforms/as_time_series.py

@@ -5,7 +5,7 @@
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class AsTimeSeries(ElementwiseTransform):
     """The `.as_time_series` transform can be used to indicate that variables shall be treated as time series.
 

bayesflow/adapters/transforms/broadcast.py

@@ -6,7 +6,7 @@
 from .transform import Transform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Broadcast(Transform):
     """
     Broadcasts arrays or scalars to the shape of a given other array.

bayesflow/adapters/transforms/concatenate.py

@@ -7,7 +7,7 @@
 from .transform import Transform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Concatenate(Transform):
     """Concatenate multiple arrays into a new key. Used to specify how data variables should be treated by the network.
 

bayesflow/adapters/transforms/constrain.py

@@ -11,7 +11,7 @@
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Constrain(ElementwiseTransform):
     """
     Constrains neural network predictions of a data variable to specified bounds.

bayesflow/adapters/transforms/convert_dtype.py

@@ -5,7 +5,7 @@
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class ConvertDType(ElementwiseTransform):
     """
     Default transform used to convert all floats from float64 to float32 to be in line with keras framework.

bayesflow/adapters/transforms/drop.py

@@ -5,7 +5,7 @@
 from .transform import Transform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Drop(Transform):
     """
     Transform to drop variables from further calculation.

bayesflow/adapters/transforms/elementwise_transform.py

@@ -3,7 +3,7 @@
 from bayesflow.utils.serialization import serializable, deserialize
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class ElementwiseTransform:
     """Base class on which other transforms are based"""
 

bayesflow/adapters/transforms/expand_dims.py

@@ -5,7 +5,7 @@
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class ExpandDims(ElementwiseTransform):
     """
     Expand the shape of an array.

bayesflow/adapters/transforms/filter_transform.py

@@ -14,7 +14,7 @@ def __call__(self, key: str, value: np.ndarray, inverse: bool) -> bool:
         raise NotImplementedError
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class FilterTransform(Transform):
     """
     Implements a transform that applies a different transform on a subset of the data.

bayesflow/adapters/transforms/keep.py

@@ -5,7 +5,7 @@
 from .transform import Transform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Keep(Transform):
     """
     Name the data parameters that should be kept for futher calculation.

bayesflow/adapters/transforms/log.py

@@ -5,7 +5,7 @@
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Log(ElementwiseTransform):
     """Log transforms a variable.
 

bayesflow/adapters/transforms/map_transform.py

@@ -6,7 +6,7 @@
 from .transform import Transform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class MapTransform(Transform):
     """
     Implements a transform that applies a set of elementwise transforms

bayesflow/adapters/transforms/numpy_transform.py

@@ -5,7 +5,7 @@
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class NumpyTransform(ElementwiseTransform):
     """
     A class to apply element-wise transformations using plain NumPy functions.

bayesflow/adapters/transforms/one_hot.py

@@ -6,7 +6,7 @@
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class OneHot(ElementwiseTransform):
     """
     Changes data to be one-hot encoded.

bayesflow/adapters/transforms/rename.py

@@ -3,7 +3,7 @@
 from .transform import Transform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Rename(Transform):
     """
     Transform to rename keys in data dictionary. Useful to rename variables to match those required by

bayesflow/adapters/transforms/scale.py

@@ -5,7 +5,7 @@
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Scale(ElementwiseTransform):
     def __init__(self, scale: np.typing.ArrayLike):
         self.scale = np.array(scale)

bayesflow/adapters/transforms/serializable_custom_transform.py

@@ -1,18 +1,17 @@
 from collections.abc import Callable
 import numpy as np
 from keras.saving import (
-    deserialize_keras_object as deserialize,
-    register_keras_serializable as serializable,
-    serialize_keras_object as serialize,
     get_registered_name,
     get_registered_object,
 )
+
+from bayesflow.utils.serialization import deserialize, serializable, serialize
 from .elementwise_transform import ElementwiseTransform
 from ...utils import filter_kwargs
 import inspect
 
 
-@serializable(package="bayesflow.adapters")
+@serializable("bayesflow.adapters")
 class SerializableCustomTransform(ElementwiseTransform):
     """
     Transforms a parameter using a pair of registered serializable forward and inverse functions.

bayesflow/adapters/transforms/shift.py

@@ -5,7 +5,7 @@
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Shift(ElementwiseTransform):
     def __init__(self, shift: np.typing.ArrayLike):
         self.shift = np.array(shift)

bayesflow/adapters/transforms/split.py

@@ -6,7 +6,7 @@
 from .transform import Transform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Split(Transform):
     """This is the effective inverse of the :py:class:`~Concatenate` Transform.
 

bayesflow/adapters/transforms/sqrt.py

@@ -5,7 +5,7 @@
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Sqrt(ElementwiseTransform):
     """Square-root transform a variable.
 

bayesflow/adapters/transforms/standardize.py

@@ -7,7 +7,7 @@
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Standardize(ElementwiseTransform):
     """
     Transform that when applied standardizes data using typical z-score standardization

bayesflow/adapters/transforms/to_array.py

@@ -7,7 +7,7 @@
 from .elementwise_transform import ElementwiseTransform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class ToArray(ElementwiseTransform):
     """
     Checks provided data for any non-arrays and converts them to numpy arrays.

bayesflow/adapters/transforms/to_dict.py

@@ -6,7 +6,7 @@
 from .transform import Transform
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class ToDict(Transform):
     """Convert non-dict batches (e.g., pandas.DataFrame) to dict batches"""
 

bayesflow/adapters/transforms/transform.py

@@ -3,7 +3,7 @@
 from bayesflow.utils.serialization import serializable, deserialize
 
 
-@serializable
+@serializable("bayesflow.adapters")
 class Transform:
     """
     Base class on which other transforms are based

bayesflow/approximators/continuous_approximator.py

@@ -13,7 +13,7 @@
 from .approximator import Approximator
 
 
-@serializable
+@serializable("bayesflow.approximators")
 class ContinuousApproximator(Approximator):
     """
     Defines a workflow for performing fast posterior or likelihood inference.

bayesflow/approximators/model_comparison_approximator.py

@@ -14,7 +14,7 @@
 from .approximator import Approximator
 
 
-@serializable
+@serializable("bayesflow.approximators")
 class ModelComparisonApproximator(Approximator):
     """
     Defines an approximator for model (simulator) comparison, where the (discrete) posterior model probabilities are

bayesflow/approximators/point_approximator.py

@@ -11,7 +11,7 @@
 from .continuous_approximator import ContinuousApproximator
 
 
-@serializable
+@serializable("bayesflow.approximators")
 class PointApproximator(ContinuousApproximator):
     """
     A workflow for fast amortized point estimation of a conditional distribution.

bayesflow/distributions/diagonal_normal.py

@@ -12,7 +12,7 @@
 from .distribution import Distribution
 
 
-@serializable
+@serializable("bayesflow.distributions")
 class DiagonalNormal(Distribution):
     """Implements a backend-agnostic diagonal Gaussian distribution."""
 

bayesflow/distributions/diagonal_student_t.py

@@ -13,7 +13,7 @@
 from .distribution import Distribution
 
 
-@serializable
+@serializable("bayesflow.distributions")
 class DiagonalStudentT(Distribution):
     """Implements a backend-agnostic diagonal Student-t distribution."""
 

bayesflow/distributions/distribution.py

@@ -5,7 +5,7 @@
 from bayesflow.utils.serialization import serializable, deserialize
 
 
-@serializable
+@serializable("bayesflow.distributions")
 class Distribution(keras.Layer):
     def __init__(self, **kwargs):
         super().__init__(**layer_kwargs(kwargs))

bayesflow/distributions/mixture.py

@@ -11,7 +11,7 @@
 from bayesflow.distributions import Distribution
 
 
-@serializable
+@serializable("bayesflow.distributions")
 class Mixture(Distribution):
     """Utility class for a backend-agnostic mixture distributions."""
 

bayesflow/experimental/cif/cif.py

@@ -1,15 +1,16 @@
 import keras
-from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Shape, Tensor
+from bayesflow.utils.serialization import serializable
 
 from bayesflow.networks.inference_network import InferenceNetwork
 from bayesflow.networks.coupling_flow import CouplingFlow
 
 from .conditional_gaussian import ConditionalGaussian
 
 
-@serializable(package="bayesflow.networks")
+# disable module check, use potential module after moving from experimental
+@serializable("bayesflow.networks", disable_module_check=True)
 class CIF(InferenceNetwork):
     """Implements a continuously indexed flow (CIF) with a `CouplingFlow`
     bijection and `ConditionalGaussian` distributions p and q. Improves on

bayesflow/experimental/cif/conditional_gaussian.py

@@ -1,13 +1,14 @@
 import keras
-from keras.saving import register_keras_serializable
 import numpy as np
 from bayesflow.networks.mlp import MLP
 
 from bayesflow.types import Shape, Tensor
 from bayesflow.utils import layer_kwargs
+from bayesflow.utils.serialization import serializable
 
 
-@register_keras_serializable(package="bayesflow.networks.cif")
+# disable module check, use potential module after moving from experimental
+@serializable("bayesflow.networks", disable_module_check=True)
 class ConditionalGaussian(keras.Layer):
     """Implements a conditional gaussian distribution with neural networks for
     the means and standard deviations respectively. Bulit in reference to [1].

bayesflow/experimental/continuous_time_consistency_model.py

@@ -22,7 +22,8 @@
 from bayesflow.networks.embeddings import FourierEmbedding
 
 
-@serializable
+# disable module check, use potential module after moving from experimental
+@serializable("bayesflow.networks", disable_module_check=True)
 class ContinuousTimeConsistencyModel(InferenceNetwork):
     """Implements an sCM (simple, stable, and scalable Consistency Model)
     with continous-time Consistency Training (CT) as described in [1].

bayesflow/experimental/free_form_flow/free_form_flow.py

@@ -19,7 +19,8 @@
 from bayesflow.networks import InferenceNetwork
 
 
-@serializable
+# disable module check, use potential module after moving from experimental
+@serializable("bayesflow.networks", disable_module_check=True)
 class FreeFormFlow(InferenceNetwork):
     """Implements a dimensionality-preserving Free-form Flow.
     Incorporates ideas from [1-2].

bayesflow/experimental/resnet/dense_resnet.py

@@ -8,7 +8,8 @@
 from .double_linear import DoubleLinear
 
 
-@serializable
+# disable module check, use potential module after moving from experimental
+@serializable("bayesflow.networks", disable_module_check=True)
 class DenseResNet(keras.Sequential):
     """
     Implements the fully-connected analogue of the ResNet architecture.

bayesflow/experimental/resnet/double_conv.py

@@ -5,7 +5,8 @@
 from bayesflow.utils.serialization import deserialize, serializable, serialize
 
 
-@serializable
+# disable module check, use potential module after moving from experimental
+@serializable("bayesflow.networks", disable_module_check=True)
 class DoubleConv(keras.Sequential):
     def __init__(
         self,