Add squeeze to adapter (#289)

bayesflow-org · Dec 20, 2024 · d367cfd · d367cfd
1 parent 5038806
commit d367cfd
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/bayesflow/adapters/adapter.py b/bayesflow/adapters/adapter.py
@@ -122,12 +122,18 @@ def as_time_series(self, keys: str | Sequence[str]):
         return self
 
     def broadcast(
-        self, keys: str | Sequence[str], *, to: str, expand: str | int | tuple = "left", exclude: int | tuple = -1
+        self,
+        keys: str | Sequence[str],
+        *,
+        to: str,
+        expand: str | int | tuple = "left",
+        exclude: int | tuple = -1,
+        squeeze: int | tuple = None,
     ):
         if isinstance(keys, str):
             keys = [keys]
 
-        transform = Broadcast(keys, to=to, expand=expand, exclude=exclude)
+        transform = Broadcast(keys, to=to, expand=expand, exclude=exclude, squeeze=squeeze)
         self.transforms.append(transform)
         return self