bayinx 0.3.15__tar.gz → 0.3.17__tar.gz

{bayinx-0.3.15 → bayinx-0.3.17}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bayinx
-Version: 0.3.15
+Version: 0.3.17
 Summary: Bayesian Inference with JAX
 License-File: LICENSE
 Requires-Python: >=3.12

{bayinx-0.3.15 → bayinx-0.3.17}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bayinx"
-version = "0.3.15"
+version = "0.3.17"
 description = "Bayesian Inference with JAX"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -19,7 +19,7 @@ build-backend = "hatchling.build"
 addopts = "-q --benchmark-min-rounds=30 --benchmark-columns=rounds,mean,median,stddev --benchmark-group-by=func"
 
 [tool.bumpversion]
-current_version = "0.3.15"
+current_version = "0.3.17"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"

{bayinx-0.3.15 → bayinx-0.3.17}/src/bayinx/core/_variational.py

@@ -175,17 +175,33 @@ class Variational(eqx.Module, Generic[M]):
         data: Any = None,
         key: Key = jr.PRNGKey(0),
     ) -> Array:
-        # Sample draws from the variational approximation
-        draws: Array = self.sample(n, key)
+        # Sample a single draw to evaluate shape of output
+        draw: Array = self.sample(1, key)[0]
+        output: Array = func(self._unflatten(draw), data)
+
+        # Allocate space for results
+        results: Array = jnp.zeros((n,) + output.shape, dtype=output.dtype)
+
+        @eqx.filter_jit
+        def body_fun(i: int, state: Tuple[Key, Array]) -> Tuple[Key, Array]:
+            # Unpack state
+            key, results = state
+
+            # Update PRNG key
+            next, key = jr.split(key)
+
+            # Draw from variational
+            draw: Array = self.sample(1, key)
 
-        # Evaluate posterior predictive
-        @jax.jit
-        @jax.vmap
-        def evaluate(draw: Array):
             # Reconstruct model
             model: M = self._unflatten(draw)
 
-            # Evaluate
-            return func(model, data)
+            # Update results with output
+            results = results.at[i].set(func(model, data))
+
+            return next, results
+
+        # Evaluate draws
+        results: Array = jax.lax.fori_loop(0, n, body_fun, (key, results))[1]
 
-        return evaluate(draws)
+        return results

{bayinx-0.3.15 → bayinx-0.3.17}/tests/test_predictive.py

@@ -10,6 +10,8 @@ from bayinx.mhx.vi import Standard
 
 class NormalDist(Model):
     x: Parameter[Array]
+    other: str = "hi"
+    other2: float = 10.0
 
     def __init__(self):
         self.x = Parameter(jnp.array([0.0, 0.0]))
@@ -33,9 +35,11 @@ def test_predictive():
 
     # Posterior predictive function
     def extract_first_x(model: NormalDist, data: Any = None) -> Array:
-        return model.x()[0]
+        return model.x()
 
     # Sample from posterior
     first_x_samples: Array = vari.posterior_predictive(extract_first_x, 1000)
 
     assert abs(first_x_samples.mean()) < 0.1 and abs(first_x_samples.var() - 1.0) < 0.1
+
+# todo: make test for data argument