unify serialization pattern

Author: LarsKue

bayesflow/adapters/transforms/as_set.py

@@ -1,8 +1,10 @@
 import numpy as np
+from keras.saving import register_keras_serializable as serializable
 
 from .elementwise_transform import ElementwiseTransform
 
 
+@serializable(package="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

@@ -1,8 +1,10 @@
 import numpy as np
+from keras.saving import register_keras_serializable as serializable
 
 from .elementwise_transform import ElementwiseTransform
 
 
+@serializable(package="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/expand_dims.py

@@ -1,13 +1,14 @@
 import numpy as np
-
 from keras.saving import (
     deserialize_keras_object as deserialize,
     serialize_keras_object as serialize,
 )
+from keras.saving import register_keras_serializable as serializable
 
 from .elementwise_transform import ElementwiseTransform
 
 
+@serializable(package="bayesflow.adapters")
 class ExpandDims(ElementwiseTransform):
     """
     Expand the shape of an array.

bayesflow/adapters/transforms/log.py

@@ -1,13 +1,14 @@
 import numpy as np
-
 from keras.saving import (
     deserialize_keras_object as deserialize,
     serialize_keras_object as serialize,
 )
+from keras.saving import register_keras_serializable as serializable
 
 from .elementwise_transform import ElementwiseTransform
 
 
+@serializable(package="bayesflow.adapters")
 class Log(ElementwiseTransform):
     """Log transforms a variable.
 

bayesflow/adapters/transforms/sqrt.py

@@ -1,8 +1,10 @@
 import numpy as np
+from keras.saving import register_keras_serializable as serializable
 
 from .elementwise_transform import ElementwiseTransform
 
 
+@serializable(package="bayesflow.adapters")
 class Sqrt(ElementwiseTransform):
     """Square-root transform a variable.