squeeze: adapt example, add comment for changing batch dims

Author: vpratz

bayesflow/adapters/adapter.py

@@ -789,7 +789,8 @@ def squeeze(self, keys: str | Sequence[str], *, axis: int | tuple):
         keys : str or Sequence of str
             The names of the variables to squeeze.
         axis : int or tuple
-            The axis to squeeze.
+            The axis to squeeze. As the number of batch dimensions might change, we advise using negative
+            numbers (i.e., indexing from the end instead of the start).
         """
         if isinstance(keys, str):
             keys = [keys]

bayesflow/adapters/transforms/squeeze.py

@@ -13,15 +13,16 @@ class Squeeze(ElementwiseTransform):
     Parameters
     ----------
     axis : int or tuple
-        The axis to squeeze.
+        The axis to squeeze. As the number of batch dimensions might change, we advise using negative
+        numbers (i.e., indexing from the end instead of the start).
 
     Examples
     --------
     shape (3, 1) array:
 
     >>> a = np.array([[1], [2], [3]])
 
-    >>> sq = bf.adapters.transforms.Squeeze(axis=1)
+    >>> sq = bf.adapters.transforms.Squeeze(axis=-1)
     >>> sq.forward(a).shape
     (3,)