Fix import

Author: stefanradev93

bayesflow/wrappers/mamba/__init__.py

@@ -1 +1 @@
-from mamba import Mamba
+from .mamba import Mamba

tests/test_wrappers/__init__.py

@@ -0,0 +1,22 @@
+import pytest
+
+
+@pytest.fixture()
+def inference_network():
+    from bayesflow.networks import CouplingFlow
+
+    return CouplingFlow(depth=2)
+
+
+@pytest.fixture()
+def random_time_series():
+    import keras
+
+    return keras.random.normal(shape=(2, 80, 2))
+
+
+@pytest.fixture()
+def mamba_summary_network():
+    from bayesflow.wrappers.mamba import Mamba
+
+    return Mamba(summary_dim=4, feature_dims=(2, 2), state_dims=(4, 4), conv_dims=(8, 8))

tests/test_wrappers/conftest.py

@@ -0,0 +1,25 @@
+import pytest
+
+import bayesflow as bf
+
+
+@pytest.mark.torch
+def test_mamba_summary(random_time_series, mamba_summary_network):
+    out = mamba_summary_network(random_time_series)
+    # Batch size 2, summary dim 4
+    assert out.shape == (2, 4)
+
+
+@pytest.mark.torch
+def test_mamba_trains(random_time_series, inference_network, mamba_summary_network):
+    workflow = bf.BasicWorkflow(
+        inference_network=inference_network,
+        summary_network=mamba_summary_network,
+        inference_variables=["parameters"],
+        summary_variables=["observables"],
+        simulator=bf.simulators.SIR(),
+    )
+
+    history = workflow.fit_online(epochs=2, batch_size=8, num_batches_per_epoch=2)
+    assert "loss" in list(history.history.keys())
+    assert len(history.history["loss"]) == 2