Add docs for samplers and improve API

Author: dirmeier

.github/workflows/ci.yaml

@@ -73,7 +73,7 @@ jobs:
           pip install hatch
       - name: Build package
         run: |
-          pip install jaxlib==0.4.24 jax==0.4.24
+          pip install jaxlib jax
       - name: Run tests
         run: |
           hatch run test:test

pyproject.toml

@@ -8,7 +8,7 @@ description = " Simulation-based inference in JAX"
 authors = [{name = "Simon Dirmeier", email = "sfyrbnd@pm.me"}]
 readme = "README.md"
 license = "Apache-2.0"
-keywords = [ "sbi", "abc", "simulation-based inference", "approximate Bayesian computation"]
+keywords = ["sbi", "abc", "simulation-based inference", "approximate Bayesian computation"]
 classifiers = [
     "Development Status :: 4 - Beta",
     "Intended Audience :: Science/Research",

sbijax/_src/npe.py

@@ -352,7 +352,7 @@ def sample_posterior(
                 proposal_probs = self.prior.log_prob(proposal)
                 proposal = jax.vmap(lambda x: ravel_pytree(x)[0])(proposal)
             else:
-                proposal_probs = self.prior_log_density_fn(
+                proposal_probs = self.prior.log_prob(
                     jax.vmap(unravel_fn)(proposal)
                 )
             if check_proposal_probs: