fix minor bugs

bayesflow-org · Jul 2, 2024 · 8b80304 · 8b80304
1 parent 8c67dc2
commit 8b80304
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/bayesflow/simulators/sequential_simulator.py b/bayesflow/simulators/sequential_simulator.py
@@ -1,15 +1,15 @@
 from collections.abc import Sequence
 import keras
 
-from .composite_simulator import Compose
+from .composite_simulator import CompositeSimulator
 from .functional_simulator import FunctionalSimulator
 from .simulator import Simulator
 from ..types import Shape, Tensor
 
 
 class SequentialSimulator(Simulator):
     def __init__(self, sample_fns: Sequence[callable], *, convert_dtypes: str = None, **kwargs):
-        self.inner = Compose([FunctionalSimulator(fn, **kwargs) for fn in sample_fns])
+        self.inner = CompositeSimulator([FunctionalSimulator(fn, **kwargs) for fn in sample_fns])
         self.convert_dtypes = convert_dtypes
 
     def sample(self, batch_shape: Shape, **kwargs) -> dict[str, Tensor]:

diff --git a/bayesflow/simulators/two_moons_simulator.py b/bayesflow/simulators/two_moons_simulator.py
@@ -8,7 +8,7 @@
 class TwoMoonsSimulator(Simulator):
     """TODO: Docs"""
 
-    def sample(self, batch_shape: Shape) -> dict[str, Tensor]:
+    def sample(self, batch_shape: Shape, **kwargs) -> dict[str, Tensor]:
         r = keras.random.normal(batch_shape + (1,), 0.1, 0.01)
         alpha = keras.random.uniform(batch_shape + (1,), -0.5 * np.pi, 0.5 * np.pi)