Add squeeze transform

Very basic transform, just the inverse of expand_dims

Author: vpratz

bayesflow/adapters/adapter.py

@@ -22,6 +22,7 @@
     OneHot,
     Rename,
     SerializableCustomTransform,
+    Squeeze,
     Sqrt,
     Standardize,
     ToArray,
@@ -780,6 +781,23 @@ def split(self, key: str, *, into: Sequence[str], indices_or_sections: int | Seq
 
         return self
 
+    def squeeze(self, keys: str | Sequence[str], *, axis: int | tuple):
+        """Append a :py:class:`~transforms.Squeeze` transform to the adapter.
+
+        Parameters
+        ----------
+        keys : str or Sequence of str
+            The names of the variables to squeeze.
+        axis : int or tuple
+            The axis to squeeze.
+        """
+        if isinstance(keys, str):
+            keys = [keys]
+
+        transform = MapTransform({key: Squeeze(axis=axis) for key in keys})
+        self.transforms.append(transform)
+        return self
+
     def sqrt(self, keys: str | Sequence[str]):
         """Append an :py:class:`~transforms.Sqrt` transform to the adapter.
 

bayesflow/adapters/transforms/__init__.py

@@ -19,6 +19,7 @@
 from .serializable_custom_transform import SerializableCustomTransform
 from .shift import Shift
 from .split import Split
+from .squeeze import Squeeze
 from .sqrt import Sqrt
 from .standardize import Standardize
 from .to_array import ToArray

tests/test_adapters/conftest.py

@@ -21,6 +21,7 @@ def serializable_fn(x):
         .concatenate(["x1", "x2"], into="x")
         .concatenate(["y1", "y2"], into="y")
         .expand_dims(["z1"], axis=2)
+        .squeeze("z1", axis=2)
         .log("p1")
         .constrain("p2", lower=0)
         .apply(include="p2", forward="exp", inverse="log")