aspire-inference 0.1.0a12__tar.gz → 0.1.0a13__tar.gz

{aspire_inference-0.1.0a12 → aspire_inference-0.1.0a13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aspire-inference
-Version: 0.1.0a12
+Version: 0.1.0a13
 Summary: Accelerate Sequential Posterior Inference via REuse
 Author-email: "Michael J. Williams" <michaeljw1@googlemail.com>
 License: MIT
@@ -19,7 +19,8 @@ Provides-Extra: scipy
 Requires-Dist: scipy; extra == "scipy"
 Provides-Extra: jax
 Requires-Dist: jax; extra == "jax"
-Requires-Dist: jaxlib; extra == "jax"
+Requires-Dist: jaxlib<0.8.2; python_version < "3.12" and extra == "jax"
+Requires-Dist: jaxlib; python_version >= "3.12" and extra == "jax"
 Requires-Dist: flowjax; extra == "jax"
 Provides-Extra: torch
 Requires-Dist: torch; extra == "torch"
@@ -32,6 +33,7 @@ Provides-Extra: emcee
 Requires-Dist: emcee; extra == "emcee"
 Provides-Extra: blackjax
 Requires-Dist: blackjax; extra == "blackjax"
+Requires-Dist: fastprogress<1.1.0; extra == "blackjax"
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-requires; extra == "test"

{aspire_inference-0.1.0a12 → aspire_inference-0.1.0a13}/aspire_inference.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aspire-inference
-Version: 0.1.0a12
+Version: 0.1.0a13
 Summary: Accelerate Sequential Posterior Inference via REuse
 Author-email: "Michael J. Williams" <michaeljw1@googlemail.com>
 License: MIT
@@ -19,7 +19,8 @@ Provides-Extra: scipy
 Requires-Dist: scipy; extra == "scipy"
 Provides-Extra: jax
 Requires-Dist: jax; extra == "jax"
-Requires-Dist: jaxlib; extra == "jax"
+Requires-Dist: jaxlib<0.8.2; python_version < "3.12" and extra == "jax"
+Requires-Dist: jaxlib; python_version >= "3.12" and extra == "jax"
 Requires-Dist: flowjax; extra == "jax"
 Provides-Extra: torch
 Requires-Dist: torch; extra == "torch"
@@ -32,6 +33,7 @@ Provides-Extra: emcee
 Requires-Dist: emcee; extra == "emcee"
 Provides-Extra: blackjax
 Requires-Dist: blackjax; extra == "blackjax"
+Requires-Dist: fastprogress<1.1.0; extra == "blackjax"
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-requires; extra == "test"

{aspire_inference-0.1.0a12 → aspire_inference-0.1.0a13}/aspire_inference.egg-info/requires.txt

@@ -6,15 +6,21 @@ h5py
 
 [blackjax]
 blackjax
+fastprogress<1.1.0
 
 [emcee]
 emcee
 
 [jax]
 jax
-jaxlib
 flowjax
 
+[jax:python_version < "3.12"]
+jaxlib<0.8.2
+
+[jax:python_version >= "3.12"]
+jaxlib
+
 [minipcn]
 minipcn[array-api]>=0.2.0a3
 orng

{aspire_inference-0.1.0a12 → aspire_inference-0.1.0a13}/pyproject.toml

@@ -30,8 +30,9 @@ scipy = [
 ]
 jax = [
     "jax",
-    "jaxlib",
-    "flowjax"
+    "jaxlib<0.8.2; python_version < '3.12'",
+    "jaxlib; python_version >= '3.12'",
+    "flowjax",
 ]
 torch = [
     "torch",
@@ -47,6 +48,7 @@ emcee = [
 ]
 blackjax = [
     "blackjax",
+    "fastprogress<1.1.0",
 ]
 test = [
     "pytest",

{aspire_inference-0.1.0a12 → aspire_inference-0.1.0a13}/src/aspire/aspire.py

@@ -354,6 +354,7 @@ class Aspire:
             xp=self.xp,
             dtype=self.dtype,
             preconditioning_transform=transform,
+            parameters=self.parameters,
             **kwargs,
         )
         return sampler

{aspire_inference-0.1.0a12 → aspire_inference-0.1.0a13}/src/aspire/samplers/mcmc.py

@@ -13,14 +13,13 @@ class MCMCSampler(Sampler):
         n_samples_drawn = 0
         samples = None
         while n_samples_drawn < n_samples:
-            n_to_draw = n_samples - n_samples_drawn
-            x, log_q = self.prior_flow.sample_and_log_prob(n_to_draw)
+            x, log_q = self.prior_flow.sample_and_log_prob(n_samples)
             new_samples = Samples(x, xp=self.xp, log_q=log_q, dtype=self.dtype)
             new_samples.log_prior = new_samples.array_to_namespace(
                 self.log_prior(new_samples)
             )
             valid = self.xp.isfinite(new_samples.log_prior)
-            n_valid = self.xp.sum(valid)
+            n_valid = int(self.xp.sum(valid))
             if n_valid > 0:
                 if samples is None:
                     samples = new_samples[valid]

{aspire_inference-0.1.0a12 → aspire_inference-0.1.0a13}/src/aspire/samplers/smc/base.py

@@ -95,6 +95,7 @@ class SMCSampler(MCMCSampler):
         beta: float,
         beta_step: float,
         min_step: float,
+        beta_tolerance: float = 1e-6,
     ) -> tuple[float, float]:
         """Determine the next beta value.
 
@@ -108,6 +109,8 @@ class SMCSampler(MCMCSampler):
             The fixed beta step size if not adaptive.
         min_step : float
             The minimum beta step size.
+        beta_tolerance : float
+            Tolerance when checking for beta convergence.
 
         Returns
         -------
@@ -124,14 +127,13 @@ class SMCSampler(MCMCSampler):
             beta_prev = beta
             beta_min = beta_prev
             beta_max = 1.0
-            tol = 1e-5
             eff_beta_max = effective_sample_size(
                 samples.log_weights(beta_max)
             ) / len(samples)
             if eff_beta_max >= self.current_target_efficiency(beta_prev):
                 beta_min = 1.0
             target_eff = self.current_target_efficiency(beta_prev)
-            while beta_max - beta_min > tol:
+            while beta_max - beta_min > beta_tolerance:
                 beta_try = 0.5 * (beta_max + beta_min)
                 eff = effective_sample_size(
                     samples.log_weights(beta_try)
@@ -163,6 +165,7 @@ class SMCSampler(MCMCSampler):
         checkpoint_every: int | None = None,
         checkpoint_file_path: str | None = None,
         resume_from: str | bytes | dict | None = None,
+        beta_tolerance: float = 1e-6,
     ) -> SMCSamples:
         resumed = resume_from is not None
         if resumed:
@@ -249,6 +252,7 @@ class SMCSampler(MCMCSampler):
                     beta,
                     beta_step,
                     min_step,
+                    beta_tolerance=beta_tolerance,
                 )
                 self.history.eff_target.append(
                     self.current_target_efficiency(beta)