update serialization for mamba

Author: LarsKue

bayesflow/wrappers/mamba/mamba.py

@@ -1,17 +1,15 @@
 from collections.abc import Sequence
 
 import keras
-from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.networks.summary_network import SummaryNetwork
 from bayesflow.types import Tensor
-from bayesflow.utils import keras_kwargs
-from bayesflow.utils.decorators import sanitize_input_shape
+from bayesflow.utils.serialization import serializable
 
 from .mamba_block import MambaBlock
 
 
-@serializable("bayesflow.wrappers")
+@serializable
 class Mamba(SummaryNetwork):
     """
     Wraps a sequence of Mamba modules using the simple Mamba module from:
@@ -71,7 +69,7 @@ def __init__(
             Additional keyword arguments passed to the `SummaryNetwork` parent class.
         """
 
-        super().__init__(**keras_kwargs(kwargs))
+        super().__init__(**kwargs)
 
         if device != "cuda":
             raise NotImplementedError("MambaSSM only supports cuda as `device`.")
@@ -115,8 +113,3 @@ def call(self, time_series: Tensor, training: bool = True, **kwargs) -> Tensor:
         summary = self.summary_stats(summary)
 
         return summary
-
-    @sanitize_input_shape
-    def build(self, input_shape):
-        super().build(input_shape)
-        self.call(keras.ops.zeros(input_shape))

bayesflow/wrappers/mamba/mamba_block.py

@@ -1,12 +1,12 @@
 import keras
-from keras.saving import register_keras_serializable as serializable
 
 from bayesflow.types import Tensor
-from bayesflow.utils import keras_kwargs
+from bayesflow.utils import layer_kwargs
 from bayesflow.utils.decorators import sanitize_input_shape
+from bayesflow.utils.serialization import serializable
 
 
-@serializable("bayesflow.wrappers")
+@serializable
 class MambaBlock(keras.Layer):
     """
     Wraps the original Mamba module from, with added functionality for bidirectional processing:
@@ -53,7 +53,7 @@ def __init__(
             Additional keyword arguments passed to the `keras.layers.Layer` initializer.
         """
 
-        super().__init__(**keras_kwargs(kwargs))
+        super().__init__(**layer_kwargs(kwargs))
 
         if keras.backend.backend() != "torch":
             raise RuntimeError("Mamba is only available using torch backend.")