custom transform: less verbose naming, fix tests

Author: vpratz

bayesflow/adapters/adapter.py

@@ -288,8 +288,8 @@ def apply_serializable(
         self,
         include: str | Sequence[str] = None,
         *,
-        serializable_forward_fn: Callable[[np.ndarray, ...], np.ndarray],
-        serializable_inverse_fn: Callable[[np.ndarray, ...], np.ndarray],
+        forward: Callable[[np.ndarray, ...], np.ndarray],
+        inverse: Callable[[np.ndarray, ...], np.ndarray],
         predicate: Predicate = None,
         exclude: str | Sequence[str] = None,
         **kwargs,
@@ -298,12 +298,12 @@ def apply_serializable(
 
         Parameters
         ----------
-        serializable_forward_fn : function, no lambda
+        forward : function, no lambda
             Registered serializable function to transform the data in the forward pass.
             For the adapter to be serializable, this function has to be serializable
             as well (see Notes). Therefore, only proper functions and no lambda
             functions can be used here.
-        serializable_inverse_fn : function, no lambda
+        inverse : function, no lambda
             Registered serializable function to transform the data in the inverse pass.
             For the adapter to be serializable, this function has to be serializable
             as well (see Notes). Therefore, only proper functions and no lambda
@@ -350,17 +350,17 @@ def apply_serializable(
         >>>
         >>> adapter = bf.Adapter().apply_serializable(
         >>>     "x",
-        >>>     serializable_forward_fn=forward_fn,
-        >>>     serializable_inverse_fn=inverse_fn,
+        >>>     forward=forward_fn,
+        >>>     inverse=inverse_fn,
         >>> )
         """
         transform = FilterTransform(
             transform_constructor=SerializableCustomTransform,
             predicate=predicate,
             include=include,
             exclude=exclude,
-            serializable_forward_fn=serializable_forward_fn,
-            serializable_inverse_fn=serializable_inverse_fn,
+            forward=forward,
+            inverse=inverse,
             **kwargs,
         )
         self.transforms.append(transform)

bayesflow/adapters/transforms/serializable_custom_transform.py

@@ -19,12 +19,12 @@ class SerializableCustomTransform(ElementwiseTransform):
 
     Parameters
     ----------
-    serializable_forward_fn : function, no lambda
+    forward : function, no lambda
         Registered serializable function to transform the data in the forward pass.
         For the adapter to be serializable, this function has to be serializable
         as well (see Notes). Therefore, only proper functions and no lambda
         functions can be used here.
-    serializable_inverse_fn : function, no lambda
+    inverse : function, no lambda
         Function to transform the data in the inverse pass.
         For the adapter to be serializable, this function has to be serializable
         as well (see Notes). Therefore, only proper functions and no lambda
@@ -47,28 +47,27 @@ class SerializableCustomTransform(ElementwiseTransform):
     def __init__(
         self,
         *,
-        serializable_forward_fn: Callable[[np.ndarray, ...], np.ndarray],
-        serializable_inverse_fn: Callable[[np.ndarray, ...], np.ndarray],
+        forward: Callable[[np.ndarray, ...], np.ndarray],
+        inverse: Callable[[np.ndarray, ...], np.ndarray],
     ):
         super().__init__()
 
-        self._check_serializable(serializable_forward_fn, label="serializable_forward_fn")
-        self._check_serializable(serializable_inverse_fn, label="serializable_inverse_fn")
-        self._forward = serializable_forward_fn
-        self._inverse = serializable_inverse_fn
+        self._check_serializable(forward, label="forward")
+        self._check_serializable(inverse, label="inverse")
+        self._forward = forward
+        self._inverse = inverse
 
     @classmethod
     def _check_serializable(cls, function, label=""):
-        GENERAL_EXAMPLE_CODE = f"""The example code below shows the structure of a correctly decorated function:
-
-```
-import keras
-
-@keras.saving.register_keras_serializable('custom')
-def my_{label}(...):
-    [your code goes here...]
-```
-"""
+        GENERAL_EXAMPLE_CODE = (
+            "The example code below shows the structure of a correctly decorated function:\n\n"
+            "```\n"
+            "import keras\n\n"
+            "@keras.saving.register_keras_serializable('custom')\n"
+            f"def my_{label}(...):\n"
+            "    [your code goes here...]\n"
+            "```\n"
+        )
         if function is None:
             raise TypeError(
                 f"'{label}' must be a registered serializable function, was 'NoneType'.\n{GENERAL_EXAMPLE_CODE}"
@@ -132,7 +131,7 @@ def from_config(cls, config: dict, custom_objects=None) -> "SerializableCustomTr
             raise TypeError(
                 "\n\nPLEASE READ HERE:\n"
                 "-----------------\n"
-                "The forward function that was provided as `serializable_forward_fn` "
+                "The forward function that was provided as `forward` "
                 "is not registered with Keras, making deserialization impossible. "
                 f"Please ensure that it is registered as '{config['forward']['config']}' and identical to the original "
                 "function before loading your model."
@@ -147,7 +146,7 @@ def from_config(cls, config: dict, custom_objects=None) -> "SerializableCustomTr
             raise TypeError(
                 "\n\nPLEASE READ HERE:\n"
                 "-----------------\n"
-                "The inverse function that was provided as `serializable_inverse_fn` "
+                "The inverse function that was provided as `inverse` "
                 "is not registered with Keras, making deserialization impossible. "
                 f"Please ensure that it is registered as '{config['inverse']['config']}' and identical to the original "
                 "function before loading your model."
@@ -156,8 +155,8 @@ def from_config(cls, config: dict, custom_objects=None) -> "SerializableCustomTr
         forward = deserialize(config["forward"], custom_objects)
         inverse = deserialize(config["inverse"], custom_objects)
         return cls(
-            serializable_forward_fn=forward,
-            serializable_inverse_fn=inverse,
+            forward=forward,
+            inverse=inverse,
         )
 
     def get_config(self) -> dict:

tests/test_adapters/conftest.py

@@ -25,9 +25,7 @@ def serializable_fn(x):
         .constrain("p2", lower=0)
         .apply(include="p2", forward="exp", inverse="log")
         .apply(include="p2", forward="log1p")
-        .apply_serializable(
-            include="x", serializable_forward_fn=serializable_fn, serializable_inverse_fn=serializable_fn
-        )
+        .apply_serializable(include="x", forward=serializable_fn, inverse=serializable_fn)
         .scale("x", by=[-1, 2])
         .shift("x", by=2)
         .standardize(exclude=["t1", "t2", "o1"])

tests/test_adapters/test_adapters.py

@@ -140,31 +140,29 @@ def registered_but_changed(x):  # noqa: F811
 
     # method instead of function provided
     with pytest.raises(ValueError):
-        SerializableCustomTransform(serializable_forward_fn=A.fn, serializable_inverse_fn=registered_fn)
-        SerializableCustomTransform(serializable_forward_fn=registered_fn, serializable_inverse_fn=A.fn)
+        SerializableCustomTransform(forward=A.fn, inverse=registered_fn)
+    with pytest.raises(ValueError):
+        SerializableCustomTransform(forward=registered_fn, inverse=A.fn)
 
     # lambda function provided
     with pytest.raises(ValueError):
-        SerializableCustomTransform(serializable_forward_fn=lambda x: x, serializable_inverse_fn=registered_fn)
-        SerializableCustomTransform(serializable_forward_fn=registered_fn, serializable_inverse_fn=lambda x: x)
+        SerializableCustomTransform(forward=lambda x: x, inverse=registered_fn)
+    with pytest.raises(ValueError):
+        SerializableCustomTransform(forward=registered_fn, inverse=lambda x: x)
 
     # unregistered function provided
     with pytest.raises(ValueError):
-        SerializableCustomTransform(serializable_forward_fn=not_registered_fn, serializable_inverse_fn=registered_fn)
-        SerializableCustomTransform(serializable_forward_fn=registered_fn, serializable_inverse_fn=not_registered_fn)
+        SerializableCustomTransform(forward=not_registered_fn, inverse=registered_fn)
+    with pytest.raises(ValueError):
+        SerializableCustomTransform(forward=registered_fn, inverse=not_registered_fn)
 
     # function does not match registered function
     with pytest.raises(ValueError):
-        SerializableCustomTransform(
-            serializable_forward_fn=registered_but_changed, serializable_inverse_fn=registered_fn
-        )
-        SerializableCustomTransform(
-            serializable_forward_fn=registered_fn, serializable_inverse_fn=registered_but_changed
-        )
-
-    transform = SerializableCustomTransform(
-        serializable_forward_fn=registered_fn, serializable_inverse_fn=registered_fn
-    )
+        SerializableCustomTransform(forward=registered_but_changed, inverse=registered_fn)
+    with pytest.raises(ValueError):
+        SerializableCustomTransform(forward=registered_fn, inverse=registered_but_changed)
+
+    transform = SerializableCustomTransform(forward=registered_fn, inverse=registered_fn)
     serialized_transform = keras.saving.serialize_keras_object(transform)
     keras.saving.deserialize_keras_object(serialized_transform)