clean up take transform and docs

Author: LarsKue

bayesflow/adapters/adapter.py

@@ -765,7 +765,7 @@ def standardize(
             Names of variables to include in the transform.
         exclude : str or Sequence of str, optional
             Names of variables to exclude from the transform.
-        **kwargs : dict
+        **kwargs :
             Additional keyword arguments passed to the transform.
         """
         transform = FilterTransform(
@@ -780,33 +780,36 @@ def standardize(
 
     def take(
         self,
-        indices,
-        axis,
+        include: str | Sequence[str] = None,
         *,
+        indices: Sequence[int],
+        axis: int = -1,
         predicate: Predicate = None,
-        include: str | Sequence[str] = None,
         exclude: str | Sequence[str] = None,
-        **kwargs,
     ):
         """
         Append a :py:class:`~transforms.Take` transform to the adapter.
 
         Parameters
         ----------
-        predicate : Predicate, optional
-            Function that indicates which variables should be transformed.
         include : str or Sequence of str, optional
             Names of variables to include in the transform.
+        indices : Sequence of int
+            Which indices to take from the data.
+        axis : int, optional
+            Which axis to take from. The last axis is used by default.
+        predicate : Predicate, optional
+            Function that indicates which variables should be transformed.
         exclude : str or Sequence of str, optional
             Names of variables to exclude from the transform.
-        **kwargs : dict
-            Additional keyword arguments passed to the transform."""
+        """
         transform = FilterTransform(
-            transform_constructor=Take(indices=indices, axis=axis),
+            transform_constructor=Take,
             predicate=predicate,
             include=include,
             exclude=exclude,
-            **kwargs,
+            indices=indices,
+            axis=axis,
         )
         self.transforms.append(transform)
         return self

bayesflow/adapters/transforms/take.py

@@ -1,5 +1,7 @@
+from collections.abc import Sequence
 import numpy as np
-from keras.saving import register_keras_serializable as serializable
+
+from bayesflow.utils.serialization import serializable
 
 from .elementwise_transform import ElementwiseTransform
 
@@ -8,19 +10,17 @@
 class Take(ElementwiseTransform):
     """
     A transform to reduce the dimensionality of arrays output by the summary network
-    Axis is a mandatory argument and will default  to the last axis.
     Example: adapter.take("x", np.arange(0,3), axis=-1)
-
     """
 
-    def __init__(self, indices, axis=-1):
+    def __init__(self, indices: Sequence[int], axis: int = -1):
         super().__init__()
         self.indices = indices
         self.axis = axis
 
-    def forward(self, data):
+    def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
         return np.take(data, self.indices, self.axis)
 
-    def inverse(self, data):
+    def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
         # not a true invertible function
         return data

tests/test_adapters/conftest.py

@@ -32,7 +32,7 @@ def serializable_fn(x):
         .standardize(exclude=["t1", "t2", "o1"])
         .drop("d1")
         .one_hot("o1", 10)
-        .keep(["x", "y", "z1", "p1", "p2", "s1", "s2", "t1", "t2", "o1", "split_1", "split_2"])
+        .keep(["x", "y", "z1", "p1", "p2", "s1", "s2", "s3", "t1", "t2", "o1", "split_1", "split_2"])
         .rename("o1", "o2")
         .random_subsample("s3", sample_size=33, axis=0)
         .take("s3", indices=np.arange(0, 32), axis=0)