fix split and tests

Author: LarsKue

bayesflow/adapters/transforms/split.py

@@ -1,9 +1,16 @@
 from collections.abc import Sequence
 import numpy as np
 
+from keras.saving import (
+    deserialize_keras_object as deserialize,
+    register_keras_serializable as serializable,
+    serialize_keras_object as serialize,
+)
+
 from .transform import Transform
 
 
+@serializable(package="bayesflow.adapters")
 class Split(Transform):
     """This is the effective inverse of the :py:class:`~Concatenate` Transform.
 
@@ -29,6 +36,23 @@ def __init__(self, key: str, into: Sequence[str], indices_or_sections: int | Seq
 
         self.indices_or_sections = indices_or_sections
 
+    @classmethod
+    def from_config(cls, config: dict, custom_objects=None) -> "Split":
+        return cls(
+            key=deserialize(config["key"], custom_objects),
+            into=deserialize(config["into"], custom_objects),
+            indices_or_sections=deserialize(config["indices_or_sections"], custom_objects),
+            axis=deserialize(config["axis"], custom_objects),
+        )
+
+    def get_config(self) -> dict:
+        return {
+            "key": serialize(self.key),
+            "into": serialize(self.into),
+            "indices_or_sections": serialize(self.indices_or_sections),
+            "axis": serialize(self.axis),
+        }
+
     def forward(self, data: dict[str, np.ndarray], strict: bool = True, **kwargs) -> dict[str, np.ndarray]:
         # avoid side effects
         data = data.copy()
@@ -39,7 +63,7 @@ def forward(self, data: dict[str, np.ndarray], strict: bool = True, **kwargs) ->
             # we cannot produce a result, but also don't have to
             return data
 
-        splits = np.split(data.pop(self.key), self.indices_or_sections)
+        splits = np.split(data.pop(self.key), self.indices_or_sections, axis=self.axis)
 
         if len(splits) != len(self.into):
             raise ValueError(f"Requested {len(self.into)} splits, but produced {len(splits)}.")

tests/test_adapters/conftest.py

@@ -28,12 +28,11 @@ def serializable_fn(x):
         .apply_serializable(include="x", forward=serializable_fn, inverse=serializable_fn)
         .scale("x", by=[-1, 2])
         .shift("x", by=2)
-        .split("x", into=["x1", "x2"])
-        .concatenate(["x1", "x2"], into="x")
+        .split("key_to_split", into=["split_1", "split_2"])
         .standardize(exclude=["t1", "t2", "o1"])
         .drop("d1")
         .one_hot("o1", 10)
-        .keep(["x", "y", "z1", "p1", "p2", "s1", "s2", "t1", "t2", "o1"])
+        .keep(["x", "y", "z1", "p1", "p2", "s1", "s2", "t1", "t2", "o1", "split_1", "split_2"])
         .rename("o1", "o2")
     )
 
@@ -57,4 +56,5 @@ def random_data():
         "d1": np.random.standard_normal(size=(32, 2)),
         "d2": np.random.standard_normal(size=(32, 2)),
         "o1": np.random.randint(0, 9, size=(32, 2)),
+        "key_to_split": np.random.standard_normal(size=(32, 10)),
     }

tests/test_adapters/test_adapters.py

@@ -177,3 +177,18 @@ def registered_but_changed(x):  # noqa: F811
     corrupt_serialized_transform["config"]["inverse"]["config"] = "nonexistent"
     with pytest.raises(TypeError):
         keras.saving.deserialize_keras_object(corrupt_serialized_transform)
+
+
+def test_split_transform(adapter, random_data):
+    assert "key_to_split" in random_data
+
+    shape = random_data["key_to_split"].shape
+    target_shape = (*shape[:-1], shape[-1] // 2)
+
+    processed = adapter(random_data)
+
+    assert "split_1" in processed
+    assert processed["split_1"].shape == target_shape
+
+    assert "split_2" in processed
+    assert processed["split_2"].shape == target_shape