Simplify and improve clamp

bayesflow/networks/coupling_flow/transforms/affine_transform.py

@@ -1,4 +1,3 @@
-import keras
 import keras.ops as ops
 from keras.saving import register_keras_serializable as serializable
 
@@ -9,19 +8,9 @@
 
 @serializable(package="networks.coupling_flow")
 class AffineTransform(Transform):
-    def __init__(self, clamp: bool | int | float | None = 3.0, **kwargs):
+    def __init__(self, clamp: bool = True, **kwargs):
         super().__init__(**kwargs)
-        match clamp:
-            case True:
-                self.clamp_factor = 3.0
-            case False:
-                self.clamp_factor = None
-            case int() | float():
-                self.clamp_factor = float(clamp)
-            case None:
-                self.clamp_factor = None
-            case _:
-                raise ValueError(f"Invalid value for 'clamp': {clamp}")
+        self.clamp = clamp
 
     @property
     def params_per_dim(self):
@@ -39,8 +28,8 @@ def constrain_parameters(self, parameters: dict[str, Tensor]) -> dict[str, Tenso
         scale = shifted_softplus(scale)
 
         # soft clamp
-        if self.clamp_factor is not None:
-            scale = self.clamp_factor * keras.ops.tanh(scale)
+        if self.clamp:
+            scale = ops.arcsinh(scale)
 
         parameters["scale"] = scale
         return parameters