Move to ruff

pyproject.toml

@@ -63,8 +63,8 @@ dependencies = [
 ]
 
 [tool.hatch.envs.test.scripts]
-lint = 'ruff check surjectors'
-test = 'pytest -v --cov=./surjectors --cov-report=xml surjectors'
+lint = 'ruff check sbijax'
+test = 'pytest -v --cov=./sbijax --cov-report=xml sbijax'
 
 [tool.bandit]
 skips = ["B101"]

sbijax/_src/mcmc/__init__.py

@@ -0,0 +1,15 @@
+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.slice import sample_with_slice
+
+__all__ = [
+ "mcmc_diagnostics",
+ "sample_with_slice",
+ "sample_with_nuts",
+ "sample_with_mala",
+ "sample_with_rmh",
+ "sample_with_imh",
+]

sbijax/_src/mcmc/slice.py

@@ -20,7 +20,7 @@ def sample_with_slice(
  r"""Sample from a distribution using the No-U-Turn sampler.
 
  Args:
- rng_seq: a hk.PRNGSequence
+ rng_key: a jax random key
  lp: the logdensity you wish to sample from
  prior: a function that returns a prior sample
  n_chains: number of chains to sample

sbijax/_src/nn/make_flows.py

@@ -12,8 +12,7 @@
  Permutation,
  TransformedDistribution,
 )
-from surjectors._src.conditioners.mlp import make_mlp
-from surjectors._src.conditioners.nn.made import MADE
+from surjectors.nn import MADE, make_mlp
 from surjectors.util import unstack