add extra sample size to simulator rejection_sample

stefanradev93 · Jul 9, 2024 · 00a3725 · 00a3725
1 parent da7bfc3
commit 00a3725
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 60 deletions.
diff --git a/bayesflow/simulators/simulator.py b/bayesflow/simulators/simulator.py
@@ -9,12 +9,27 @@ def sample(self, batch_shape: Shape, *, numpy: bool = False, **kwargs) -> dict[s
  raise NotImplementedError
 
  def rejection_sample(
- self, batch_shape: Shape, condition: callable, *, axis: int = 0, numpy: bool = False, **kwargs
+ self,
+ batch_shape: Shape,
+ condition: callable,
+ *,
+ axis: int = 0,
+ numpy: bool = False,
+ extra_sample_size: int = None,
+ **kwargs,
  ) -> dict[str, Tensor]:
+ if extra_sample_size is None:
+ sample_shape = batch_shape
+ else:
+ sample_shape = list(batch_shape)
+ sample_shape[axis] = extra_sample_size
+
+ sample_shape = tuple(sample_shape)
+
  result = {}
 
  while not result or keras.ops.shape(next(iter(result.values())))[axis] < batch_shape[axis]:
- samples = self.sample(batch_shape, **kwargs)
+ samples = self.sample(sample_shape, **kwargs)
  accept_mask = condition(samples)
  accept_mask = keras.ops.convert_to_tensor(accept_mask, dtype="bool")