DRAFT: Move to TF iterators (#23)

* Move all classes to use TF iterators
* Update unit tests

examples/bivariate_gaussian_snl.py

@@ -22,7 +22,7 @@
 from surjectors.util import unstack
 
 from sbijax import SNL
-from sbijax.mcmc import sample_with_slice
+from sbijax.mcmc import sample_with_nuts
 
 
 def prior_model_fns():
@@ -84,9 +84,6 @@ def _flow(method, **kwargs):
 def run():
  y_observed = jnp.array([-2.0, 1.0])
 
- log_density_partial = partial(log_density_fn, y=y_observed)
- log_density = lambda x: jax.vmap(log_density_partial)(x)
-
  prior_simulator_fn, prior_logdensity_fn = prior_model_fns()
  fns = (prior_simulator_fn, prior_logdensity_fn), simulator_fn
 
@@ -107,7 +104,9 @@ def run():
  )
 
  sample_key, rng_key = jr.split(jr.PRNGKey(123))
- slice_samples = sample_with_slice(
+ log_density_partial = partial(log_density_fn, y=y_observed)
+ log_density = lambda x: log_density_partial(**x)
+ slice_samples = sample_with_nuts(
  sample_key, log_density, prior_simulator_fn
  )
  slice_samples = slice_samples.reshape(-1, 2)

examples/bivariate_gaussian_snr.py

@@ -43,7 +43,7 @@ def run():
  optimizer = optax.adam(1e-3)
 
  data, params = None, {}
- for i in range(5):
+ for i in range(2):
  data, _ = snr.simulate_data_and_possibly_append(
  jr.fold_in(jr.PRNGKey(1), i),
  params=params,
@@ -54,7 +54,6 @@ def run():
  jr.fold_in(jr.PRNGKey(2), i),
  data=data,
  optimizer=optimizer,
- batch_size=100,
  )
 
  rng_key = jr.PRNGKey(23)

examples/slcp_snass.py

@@ -20,8 +20,8 @@
 from surjectors.nn import MADE
 from surjectors.util import unstack
 
-from sbijax import SNASSS
-from sbijax.nn import make_snasss_net
+from sbijax import SNASS
+from sbijax._src.nn.make_snass_networks import make_snass_net
 
 
 def prior_model_fns():
@@ -125,15 +125,15 @@ def run():
  prior_simulator_fn, prior_logdensity_fn = prior_model_fns()
  fns = (prior_simulator_fn, prior_logdensity_fn), simulator_fn
 
- estim = SNASSS(
+ estim = SNASS(
  fns,
  make_model(5),
- make_snasss_net([64, 64, 5], [64, 64, 1], [64, 64, 1]),
+ make_snass_net([64, 64, 5], [64, 64, 1]),
  )
  optimizer = optax.adam(1e-3)
 
  data, params = None, {}
- for i in range(5):
+ for i in range(2):
  data, _ = estim.simulate_data_and_possibly_append(
  jr.fold_in(jr.PRNGKey(12), i),
  params=params,
@@ -144,7 +144,6 @@ def run():
  jr.fold_in(jr.PRNGKey(23), i),
  data=data,
  optimizer=optimizer,
- batch_size=100,
  )
 
  rng_key = jr.PRNGKey(23)

examples/slcp_ssnl.py

@@ -15,7 +15,6 @@
 from jax import numpy as jnp
 from jax import random as jr
 from jax import scipy as jsp
-from jax import vmap
 from surjectors import (
  AffineMaskedAutoregressiveInferenceFunnel,
  Chain,
@@ -27,7 +26,7 @@
 from surjectors.util import unstack
 
 from sbijax import SNL
-from sbijax.mcmc import sample_with_slice
+from sbijax.mcmc import sample_with_nuts
 
 
 def prior_model_fns():
@@ -160,50 +159,36 @@ def _flow(method, **kwargs):
 
 def run(use_surjectors):
  len_theta = 5
- # this is the thetas used in SNL
- # thetas = jnp.array([-0.7, -2.9, -1.0, -0.9, 0.6])
- y_observed = jnp.array(
- [
- [
- -0.9707123,
- -2.9461224,
- -0.4494722,
- -3.4231849,
- -0.13285634,
- -3.364017,
- -0.85367596,
- -2.4271638,
- ]
- ]
+ thetas = jnp.linspace(-2.0, 2.0, len_theta)
+ y_0 = simulator_fn(jr.PRNGKey(0), thetas.reshape(-1, len_theta)).reshape(
+ -1, 8
  )
 
- log_density_partial = partial(log_density_fn, y=y_observed)
- log_density = lambda x: vmap(log_density_partial)(x)
-
  prior_simulator_fn, prior_fn = prior_model_fns()
  fns = (prior_simulator_fn, prior_fn), simulator_fn
 
- snl = SNL(fns, make_model(y_observed.shape[1], use_surjectors))
+ snl = SNL(fns, make_model(y_0.shape[1], use_surjectors))
  optimizer = optax.adam(1e-3)
 
  data, params = None, {}
  for i in range(5):
  data, _ = snl.simulate_data_and_possibly_append(
  jr.fold_in(jr.PRNGKey(12), i),
  params=params,
- observable=y_observed,
+ observable=y_0,
  data=data,
- sampler="slice",
  )
  params, info = snl.fit(
  jr.fold_in(jr.PRNGKey(23), i), data=data, optimizer=optimizer
  )
 
  sample_key, rng_key = jr.split(jr.PRNGKey(123))
- snl_samples, _ = snl.sample_posterior(sample_key, params, y_observed)
+ snl_samples, _ = snl.sample_posterior(sample_key, params, y_0)
 
  sample_key, rng_key = jr.split(rng_key)
- slice_samples = sample_with_slice(
+ log_density_partial = partial(log_density_fn, y=y_0)
+ log_density = lambda x: log_density_partial(**x)
+ slice_samples = sample_with_nuts(
  sample_key,
  log_density,
  prior_simulator_fn,

pyproject.toml

@@ -27,6 +27,8 @@ dependencies = [
  "optax>=0.1.3",
  "surjectors>=0.3.0",
  "tfp-nightly>=0.20.0.dev20230404",
+ "tensorflow==2.15.0",
+ "tensorflow-datasets==4.9.3",
  "tqdm>=4.64.1"
 ]
 dynamic = ["version"]

sbijax/__init__.py

@@ -2,7 +2,7 @@
 sbijax: Simulation-based inference in JAX
 """
 
-__version__ = "0.1.7"
+__version__ = "0.1.8"
 
 
 from sbijax._src.abc.smc_abc import SMCABC

sbijax/_src/_sne_base.py

@@ -6,7 +6,8 @@
 from jax import random as jr
 
 from sbijax._src._sbi_base import SBI
-from sbijax._src.generator import as_batch_iterators, named_dataset
+from sbijax._src.util.data import stack_data
+from sbijax._src.util.dataloader import as_batch_iterators, named_dataset
 
 
 # pylint: disable=too-many-arguments,unused-argument
@@ -61,7 +62,7 @@ def simulate_data_and_possibly_append(
  if data is None:
  d_new = new_data
  else:
- d_new = self.stack_data(data, new_data)
+ d_new = stack_data(data, new_data)
  return d_new, diagnostics
 
  @abc.abstractmethod
@@ -135,25 +136,6 @@ def simulate_data(
 
  return new_data, diagnostics
 
- @staticmethod
- def stack_data(data, also_data):
- """Stack two data sets.
-
- Args:
- data: one data set
- also_data: another data set
-
- Returns:
- returns the stack of the two data sets
- """
- if data is None:
- return also_data
- if also_data is None:
- return data
- return named_dataset(
- *[jnp.vstack([a, b]) for a, b in zip(data, also_data)]
- )
-
  @staticmethod
  def as_iterators(
  rng_key, data, batch_size, percentage_data_as_validation_set

sbijax/_src/mcmc/__init__.py

@@ -1,6 +1,6 @@
+from sbijax._src.mcmc.diagnostics import mcmc_diagnostics
 from sbijax._src.mcmc.irmh import sample_with_imh
 from sbijax._src.mcmc.mala import sample_with_mala
 from sbijax._src.mcmc.nuts import sample_with_nuts
 from sbijax._src.mcmc.rmh import sample_with_rmh
-from sbijax._src.mcmc.sample import mcmc_diagnostics
 from sbijax._src.mcmc.slice import sample_with_slice