bayinx 0.3.15__py3-none-any.whl → 0.3.16__py3-none-any.whl

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)
+        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.dist-info → bayinx-0.3.16.dist-info}/METADATA

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

{bayinx-0.3.15.dist-info → bayinx-0.3.16.dist-info}/RECORD

@@ -8,7 +8,7 @@ bayinx/core/_flow.py,sha256=3q4rKvATnbUpuj4ICUd4lIxu_3z7GRDuNujVhAku1X0,2342
 bayinx/core/_model.py,sha256=FJUyYVE9e2uTFamxtSMKY_VV2stiU2QF68Wl_7EAKEU,2895
 bayinx/core/_optimization.py,sha256=mmeVUqfFARz8F7q4LRl-uEwVWzekmzh-9o7PnuvsHZk,2651
 bayinx/core/_parameter.py,sha256=r20JedTW2lY0miNNh9y6LeIVAsGX1kP_rlGxphW_jZg,1080
-bayinx/core/_variational.py,sha256=bQYiN8c4AGPt4hsNT68zN7J6o0fdsLGwVjbDyl62LnI,5639
+bayinx/core/_variational.py,sha256=MbOgWs0-mkEwIYsh1Hhn7anWkQ7fAGjkr6VJybYAuR0,6185
 bayinx/dists/__init__.py,sha256=BIrypqMnTLWK3a_zw8fYKMyuEMxP_qGsLfLeScias0o,118
 bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
 bayinx/dists/gamma2.py,sha256=HtB60LUQdPj4yDAHme2jsHNmLfrAKWsSZnDYkxAGaOI,1548
@@ -31,7 +31,7 @@ bayinx/mhx/vi/flows/fullaffine.py,sha256=11y_A0oO3bkKDSz-EQ6Sf4Ec2M7vHZxw94EdvAD
 bayinx/mhx/vi/flows/planar.py,sha256=2I2WzIskl8MRpJkK13FQE3vSF-077qo8gRed_EL1Pn8,1920
 bayinx/mhx/vi/flows/radial.py,sha256=e0GfuO-CL8SVr3YnEfsxStpyKcJlFpzMyjMp3sa38hg,2503
 bayinx/mhx/vi/flows/sylvester.py,sha256=ppK0BmS_ThvrCEhJiP_-p-kj67TQHSlU_RUZpDbIhsQ,469
-bayinx-0.3.15.dist-info/METADATA,sha256=frqF_xKQsuq69EoddoTdyRwD2VwY3oNNKJnlYvG4QJs,3087
-bayinx-0.3.15.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.3.15.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
-bayinx-0.3.15.dist-info/RECORD,,
+bayinx-0.3.16.dist-info/METADATA,sha256=9oLNZAc1YvN9FpWiuFFlztL-OkhclOYPBQFkuiHmUT8,3087
+bayinx-0.3.16.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.3.16.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
+bayinx-0.3.16.dist-info/RECORD,,