more docs

dirmeier · Feb 26, 2024 · 8d7d715 · 8d7d715
1 parent a998f42
commit 8d7d715
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 3 deletions.
diff --git a/sbijax/_src/_sne_base.py b/sbijax/_src/_sne_base.py
@@ -1,3 +1,4 @@
+import abc
 from abc import ABC
 
 import chex
@@ -63,6 +64,18 @@ def simulate_data_and_possibly_append(
  d_new = self.stack_data(data, new_data)
  return d_new, diagnostics
 
+ @abc.abstractmethod
+ def sample_posterior(self, rng_key, params, observable, *args, **kwargs):
+ """Sample from the approximate posterior.
+
+ Args:
+ rng_key: a jax random key
+ params: a pytree of neural network parameters
+ observable: a data point
+ *args: argument list
+ **kwargs: keyword arguments
+ """
+
  def simulate_data(
  self,
  rng_key,

diff --git a/sbijax/_src/snl.py b/sbijax/_src/snl.py
@@ -38,6 +38,7 @@ class SNL(SNE):
  arXiv preprint arXiv:2308.01054, 2023.
  """
 
+ # pylint: disable=useless-parent-delegation
  def __init__(self, model_fns, density_estimator):
  """Construct a SNL object.
 

diff --git a/sbijax/_src/snr.py b/sbijax/_src/snr.py
@@ -315,9 +315,9 @@ def simulate_data_and_possibly_append(
 
  def sample_posterior(
  self,
- rng_key: Array,
- params: Optional[Params],
- observable: Optional[Array],
+ rng_key,
+ params,
+ observable,
  *,
  n_chains=4,
  n_samples=2_000,