refactor: pass constructed prior to estimators, not a factory

examples/gaussian_linear-aio.py

@@ -29,7 +29,7 @@ def simulator_fn(seed, theta):
 
 def run(n_iter):
   y_observed = jnp.linspace(-2.0, 2.0, 5)
-  fns = prior_fn, simulator_fn
+  fns = prior_fn(), simulator_fn
   mask = jnp.zeros((10, 10))
   mask = mask.at[np.arange(5, 10), np.arange(5)].set(1)
   mask = mask + mask.T + jnp.eye(10)

examples/gaussian_linear-smcabc.py

@@ -40,7 +40,7 @@ def distance_fn(y_simulated, y_observed):
 def run(n_rounds):
   y_observed = jnp.array([-1.0, 1.0])
 
-  fns = prior_fn, simulator_fn
+  fns = prior_fn(), simulator_fn
 
   smc = SMCABC(fns, summary_fn, distance_fn)
   smc_samples, _ = smc.sample_posterior(

examples/mixture_model-cmpe.py

@@ -35,7 +35,7 @@ def simulator_fn(seed, theta):
 
 def run(n_iter):
   y_observed = jnp.array([-2.0, 1.0])
-  fns = prior_fn, simulator_fn
+  fns = prior_fn(), simulator_fn
   neural_network = make_cm(2, 64)
   model = CMPE(fns, neural_network)
 

examples/mixture_model-fmpe.py

@@ -33,7 +33,7 @@ def simulator_fn(seed, theta):
 
 def run(n_iter):
   y_observed = jnp.array([-2.0, 2.0])
-  fns = prior_fn, simulator_fn
+  fns = prior_fn(), simulator_fn
   neural_network = make_cnf(2)
   model = FMPE(fns, neural_network)
 

examples/mixture_model-nle.py

@@ -33,7 +33,7 @@ def simulator_fn(seed, theta):
 
 def run(use_spf, n_iter):
   y_observed = jnp.array([-2.0, 1.0])
-  fns = prior_fn, simulator_fn
+  fns = prior_fn(), simulator_fn
   neural_network = (
     make_spf(2, -5.0, 5.0, n_params=10) if use_spf else make_mdn(2, 10)
   )

examples/mixture_model-npe.py

@@ -35,7 +35,7 @@ def simulator_fn(seed, theta):
 
 def run(n_iter):
   y_observed = jnp.array([-2.0, 1.0])
-  fns = prior_fn, simulator_fn
+  fns = prior_fn(), simulator_fn
   neural_network = make_maf(2)
   model = NPE(fns, neural_network, use_event_space_bijections=False)
 

examples/mixture_model-npse.py

@@ -34,7 +34,7 @@ def simulator_fn(seed, theta):
 
 def run(n_iter):
   y_observed = jnp.array([-2.0, 2.0])
-  fns = prior_fn, simulator_fn
+  fns = prior_fn(), simulator_fn
   neural_network = make_score_model(2)
   model = NPSE(fns, neural_network)
 

examples/mixture_model-nre.py

@@ -35,7 +35,7 @@ def simulator_fn(seed, theta):
 
 def run(n_iter):
   y_observed = jnp.array([-2.0, 1.0])
-  fns = prior_fn, simulator_fn
+  fns = prior_fn(), simulator_fn
   neural_network = make_mlp()
   model = NRE(fns, neural_network)
 

examples/slcp-fmpe.py

@@ -65,7 +65,7 @@ def run(n_iter):
   n_dim_theta = 5
   n_layers, hidden_size = 5, 128
   neural_network = make_cnf(n_dim_theta, n_layers, hidden_size)
-  fns = prior_fn, simulator_fn
+  fns = prior_fn(), simulator_fn
   fmpe = FMPE(fns, neural_network)
 
   data, _ = fmpe.simulate_data(

examples/slcp-nass_nle.py

@@ -66,7 +66,7 @@ def run(n_rounds, n_iter):
       ]
     ]
   )
-  fns = prior_fn, simulator_fn
+  fns = prior_fn(), simulator_fn
   neural_network = make_nass_net(5, (64, 64))
   model_nass = NASS(fns, neural_network)
   model_nle = NLE(fns, make_maf(5))

examples/slcp-nass_smcabc.py

@@ -75,7 +75,7 @@ def run(n_rounds, n_iter):
       ]
     ]
   )
-  fns = prior_fn, simulator_fn
+  fns = prior_fn(), simulator_fn
   model_nass = NASS(fns, make_nass_net(5, (64, 64)))
 
   data, _ = model_nass.simulate_data(jr.PRNGKey(1), n_simulations=20_000)

examples/slcp-snle.py

@@ -166,7 +166,7 @@ def run(n_rounds, n_iter):
       ]
     ]
   )
-  fns = prior_fn, simulator_fn
+  fns = prior_fn(), simulator_fn
 
   neural_network = make_maf(8, n_layer_dimensions=[8, 8, 5, 5, 5])
   snl = SNLE(fns, neural_network)

sbijax/_src/_sbi_base.py

@@ -14,7 +14,7 @@ def __init__(self, model_fns):
     Args:
         model_fns: tuple
     """
-    self.prior = model_fns[0]()
+    self.prior = model_fns[0]
     self.simulator_fn = model_fns[1]
 
   @staticmethod

sbijax/_src/abc/sabc.py

@@ -426,7 +426,7 @@ class SABC(SBI):
   Examples:
       >>> from sbijax import SABC
       >>> from tensorflow_probability.substrates.jax import distributions as tfd
-      >>> prior = lambda: tfd.JointDistributionNamed(
+      >>> prior = tfd.JointDistributionNamed(
       ...     dict(theta=tfd.Normal(jnp.zeros(2), 1.0)), batch_ndims=0)
       >>> sim = lambda seed, theta: theta["theta"] + tfd.Normal(
       ...     0.0, 0.1).sample(theta["theta"].shape, seed=seed)

sbijax/_src/abc/sabc_test.py

@@ -235,7 +235,7 @@ def simulator_fn(seed, theta):
       theta["theta"].shape, seed=seed
     )
 
-  return prior_fn, simulator_fn
+  return prior_fn(), simulator_fn
 
 
 @pytest.mark.parametrize(

sbijax/_src/abc/smc_abc.py

@@ -26,8 +26,8 @@ class SMCABC(SBI):
   Implements the algorithm from :cite:t:`beaumont2009adaptive`.
 
   Args:
-      model_fns: a tuple of callables. The first element needs to be a
-          function that constructs a tfd.JointDistributionNamed, the second
+      model_fns: a tuple. The first element is a
+          tfd.JointDistributionNamed prior distribution, the second
           element is a simulator function.
       summary_fn: summary function
       distance_fn: distance function
@@ -36,7 +36,7 @@ class SMCABC(SBI):
       >>> from sbijax import SMCABC
       >>> from tensorflow_probability.substrates.jax import distributions as tfd
       ...
-      >>> prior = lambda: tfd.JointDistributionNamed(
+      >>> prior = tfd.JointDistributionNamed(
       ...     dict(theta=tfd.Normal(0.0, 1.0))
       ... )
       >>> s = lambda seed, theta: tfd.Normal(theta["theta"], 1.0).sample(seed=seed)

sbijax/_src/cmpe.py

@@ -92,8 +92,8 @@ class CMPE(FMPE):
   :cite:t:`schmitt2023con`.
 
   Args:
-     model_fns: a tuple of callables. The first element needs to be a
-          function that constructs a tfd.JointDistributionNamed, the second
+     model_fns: a tuple. The first element is a
+          tfd.JointDistributionNamed prior distribution, the second
           element is a simulator function.
       network: a consistency model
       t_min: minimal time point for ODE integration
@@ -104,7 +104,7 @@ class CMPE(FMPE):
       >>> from sbijax.nn import make_cm
       >>> from tensorflow_probability.substrates.jax import distributions as tfd
       ...
-      >>> prior = lambda: tfd.JointDistributionNamed(
+      >>> prior = tfd.JointDistributionNamed(
       ...     dict(theta=tfd.Normal(0.0, 1.0))
       ... )
       >>> s = lambda seed, theta: tfd.Normal(theta["theta"], 1.0).sample(seed=seed)

sbijax/_src/conftest.py

@@ -33,4 +33,4 @@ def simulator_fn(seed, theta):
 
 @pytest.fixture()
 def prior_simulator_tuple(request):
-  yield prior_fn, simulator_fn
+  yield prior_fn(), simulator_fn

sbijax/_src/experimental/aio.py

@@ -20,8 +20,8 @@ class AiO(FMPE):
   dependencies.
 
   Args:
-      model_fns: a tuple of callables. The first element needs to be a
-          function that constructs a tfd.JointDistributionNamed, the second
+      model_fns: a tuple. The first element is a
+          tfd.JointDistributionNamed prior distribution, the second
           element is a simulator function.
       score_estimator: a score estimator
 
@@ -30,7 +30,7 @@ class AiO(FMPE):
       >>> from sbijax.experimental.nn import make_simformer_based_score_model
       >>> from tensorflow_probability.substrates.jax import distributions as tfd
       ...
-      >>> prior = lambda: tfd.JointDistributionNamed(
+      >>> prior = tfd.JointDistributionNamed(
       ...    dict(theta=tfd.Normal(jnp.zeros(2), 1.0))
       ... )
       >>> s = lambda seed, theta: tfd.Normal(theta["theta"], 1.0).sample(seed=seed)

sbijax/_src/experimental/npse.py

@@ -13,8 +13,8 @@ class NPSE(FMPE):
   :cite:t:`sharrock2024sequential`.
 
   Args:
-      model_fns: a tuple of callables. The first element needs to be a
-          function that constructs a tfd.JointDistributionNamed, the second
+      model_fns: a tuple. The first element is a
+          tfd.JointDistributionNamed prior distribution, the second
           element is a simulator function.
       score_estimator: a score_estimator estimator
 
@@ -23,7 +23,7 @@ class NPSE(FMPE):
       >>> from sbijax.experimental.nn import make_score_model
       >>> from tensorflow_probability.substrates.jax import distributions as tfd
       ...
-      >>> prior = lambda: tfd.JointDistributionNamed(
+      >>> prior = tfd.JointDistributionNamed(
       ...    dict(theta=tfd.Normal(0.0, 1.0))
       ... )
       >>> s = lambda seed, theta: tfd.Normal(theta["theta"], 1.0).sample(seed=seed)

sbijax/_src/fmpe.py

@@ -18,8 +18,8 @@ class FMPE(NE):
   Implements the FMPE algorithm introduced in :cite:t:`wilderberger2023flow`.
 
   Args:
-      model_fns: a tuple of callables. The first element needs to be a
-          function that constructs a tfd.JointDistributionNamed, the second
+      model_fns: a tuple. The first element is a
+          tfd.JointDistributionNamed prior distribution, the second
           element is a simulator function.
       density_estimator: a continuous normalizing flow model
 
@@ -28,7 +28,7 @@ class FMPE(NE):
       >>> from sbijax.nn import make_cnf
       >>> from tensorflow_probability.substrates.jax import distributions as tfd
       ...
-      >>> prior = lambda: tfd.JointDistributionNamed(
+      >>> prior = tfd.JointDistributionNamed(
       ...     dict(theta=tfd.Normal(0.0, 1.0))
       ... )
       >>> s = lambda seed, theta: tfd.Normal(theta["theta"], 1.0).sample(seed=seed)

sbijax/_src/nass.py

@@ -37,8 +37,8 @@ class NASS(NE):
   can be used to infer posterior distributions.
 
   Args:
-      model_fns: a tuple of calalbles. The first element needs to be a
-          function that constructs a tfd.JointDistributionNamed, the second
+      model_fns: a tuple. The first element is a
+          tfd.JointDistributionNamed prior distribution, the second
           element is a simulator function.
       summary_net: a SNASSNet object
 
@@ -47,7 +47,7 @@ class NASS(NE):
       >>> from sbijax.nn import make_nass_net
       >>> from tensorflow_probability.substrates.jax import distributions as tfd
       ...
-      >>> prior = lambda: tfd.JointDistributionNamed(
+      >>> prior = tfd.JointDistributionNamed(
       ...    dict(theta=tfd.Normal(jnp.zeros(5), 1.0))
       ... )
       >>> s = lambda seed, theta: tfd.Normal(

sbijax/_src/nasss.py

@@ -57,8 +57,8 @@ class NASSS(NASS):
   can be used to infer posterior distributions.
 
   Args:
-      model_fns: a tuple of calalbles. The first element needs to be a
-          function that constructs a tfd.JointDistributionNamed, the second
+      model_fns: a tuple. The first element is a
+          tfd.JointDistributionNamed prior distribution, the second
           element is a simulator function.
       summary_net: a (neural) conditional density estimator
           to model the likelihood function of summary statistics, i.e.,
@@ -71,7 +71,7 @@ class NASSS(NASS):
       >>> from sbijax.nn import make_nasss_net
       >>> from tensorflow_probability.substrates.jax import distributions as tfd
       ...
-      >>> prior = lambda: tfd.JointDistributionNamed(
+      >>> prior = tfd.JointDistributionNamed(
       ...    dict(theta=tfd.Normal(jnp.zeros(5), 1.0))
       ... )
       >>> s = lambda seed, theta: tfd.Normal(

sbijax/_src/nle.py

@@ -22,8 +22,8 @@ class NLE(NE):
   Implements the method introduced in :cite:t:`papama2019neural`.
 
   Args:
-      model_fns: a tuple of calalbles. The first element needs to be a
-          function that constructs a tfd.JointDistributionNamed, the second
+      model_fns: a tuple. The first element is a
+          tfd.JointDistributionNamed prior distribution, the second
           element is a simulator function.
       density_estimator: a (neural) conditional density estimator
           to model the likelihood function
@@ -33,7 +33,7 @@ class NLE(NE):
       >>> from sbijax.nn import make_mdn
       >>> from tensorflow_probability.substrates.jax import distributions as tfd
       ...
-      >>> prior = lambda: tfd.JointDistributionNamed(
+      >>> prior = tfd.JointDistributionNamed(
       ...    dict(theta=tfd.Normal(0.0, 1.0))
       ... )
       >>> s = lambda seed, theta: tfd.Normal(theta["theta"], 1.0).sample(seed=seed)

sbijax/_src/npe.py

@@ -22,8 +22,8 @@ class NPE(NE):
   here we refer to it simply as NPE.
 
   Args:
-      model_fns: a tuple of calalbles. The first element needs to be a
-          function that constructs a tfd.JointDistributionNamed, the second
+      model_fns: a tuple. The first element is a
+          tfd.JointDistributionNamed prior distribution, the second
           element is a simulator function.
       density_estimator: a (neural) conditional density estimator
           to model the posterior distribution
@@ -34,7 +34,7 @@ class NPE(NE):
       >>> from sbijax.nn import make_maf
       >>> from tensorflow_probability.substrates.jax import distributions as tfd
       ...
-      >>> prior = lambda: tfd.JointDistributionNamed(
+      >>> prior = tfd.JointDistributionNamed(
       ...     dict(theta=tfd.Normal(0.0, 1.0))
       ... )
       >>> s = lambda seed, theta: tfd.Normal(theta["theta"], 1.0).sample(seed=seed)
@@ -70,7 +70,7 @@ def __init__(
     super().__init__(model_fns, density_estimator)
     self.num_atoms = num_atoms
     self.n_round = 0
-    prior = model_fns[0]()
+    prior = model_fns[0]
     # TODO(simon): check out event bijections
     if (
       hasattr(prior, "experimental_default_event_space_bijector")

sbijax/_src/nre.py

@@ -95,8 +95,8 @@ class NRE(NE):
   it as NRE.
 
   Args:
-      model_fns: a tuple of calalbles. The first element needs to be a
-          function that constructs a tfd.JointDistributionNamed, the second
+      model_fns: a tuple. The first element is a
+          tfd.JointDistributionNamed prior distribution, the second
           element is a simulator function.
       classifier: a neural network for classification
       num_classes: number of classes to classify against
@@ -107,7 +107,7 @@ class NRE(NE):
       >>> from sbijax.nn import make_resnet
       >>> from tensorflow_probability.substrates.jax import distributions as tfd
       ...
-      >>> prior = lambda: tfd.JointDistributionNamed(
+      >>> prior = tfd.JointDistributionNamed(
       ...     dict(theta=tfd.Normal(0.0, 1.0))
       ... )
       >>> s = lambda seed, theta: tfd.Normal(theta["theta"], 1.0).sample(seed=seed)

sbijax/_src/snle.py

@@ -9,8 +9,8 @@ class SNLE(NLE):
   it reduces its dimensionality using dimensionality reduction.
 
   Args:
-      model_fns: a tuple of calalbles. The first element needs to be a
-          function that constructs a tfd.JointDistributionNamed, the second
+      model_fns: a tuple. The first element is a
+          tfd.JointDistributionNamed prior distribution, the second
           element is a simulator function.
       density_estimator: a (neural) conditional density estimator
           to model the likelihood function
@@ -21,7 +21,7 @@ class SNLE(NLE):
       >>> from sbijax.nn import make_maf
       >>> from tensorflow_probability.substrates.jax import distributions as tfd
       ...
-      >>> prior = lambda: tfd.JointDistributionNamed(
+      >>> prior = tfd.JointDistributionNamed(
       ...    dict(theta=tfd.Normal(jnp.zeros(5), 1.0))
       ... )
       >>> s = lambda seed, theta: tfd.Normal(