bayinx 0.3.7__py3-none-any.whl → 0.3.9__py3-none-any.whl

bayinx/core/_variational.py

@@ -163,7 +163,7 @@ class Variational(eqx.Module, Generic[M]):
 
     @eqx.filter_jit
     def posterior_predictive(
-        self, func: Callable[[M], Array], n: int, key: Key = jr.PRNGKey(0)
+        self, func: Callable[[M, Any], Array], n: int, data: Any = None, key: Key = jr.PRNGKey(0)
     ) -> Array:
         # Sample draws from the variational approximation
         draws: Array = self.sample(n, key)
@@ -171,11 +171,11 @@ class Variational(eqx.Module, Generic[M]):
         # Evaluate posterior predictive
         @jax.jit
         @jax.vmap
-        def evaluate(draw: Array):
+        def evaluate(draw: Array, data: Any = None):
             # Reconstruct model
             model: M = self._unflatten(draw)
 
             # Evaluate
-            return func(model)
+            return func(model, data)
 
-        return evaluate(draws)
+        return evaluate(draws, data)

{bayinx-0.3.7.dist-info → bayinx-0.3.9.dist-info}/METADATA

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

{bayinx-0.3.7.dist-info → bayinx-0.3.9.dist-info}/RECORD

@@ -7,7 +7,7 @@ bayinx/core/_constraint.py,sha256=Gx07ZT66VE2y-qZCmBDm3_y0wO4xQyslZW10Lec1_lM,76
 bayinx/core/_flow.py,sha256=3q4rKvATnbUpuj4ICUd4lIxu_3z7GRDuNujVhAku1X0,2342
 bayinx/core/_model.py,sha256=FJUyYVE9e2uTFamxtSMKY_VV2stiU2QF68Wl_7EAKEU,2895
 bayinx/core/_parameter.py,sha256=r20JedTW2lY0miNNh9y6LeIVAsGX1kP_rlGxphW_jZg,1080
-bayinx/core/_variational.py,sha256=b7xlUcw8JDDBfXgDLMcjsOMHpFZ2Tg3sEt965eWmctI,5431
+bayinx/core/_variational.py,sha256=X8o81b8vyU7vJxw8pZYH_nxc3u990tRUZgRhMNodNI4,5484
 bayinx/dists/__init__.py,sha256=9DdPea7HAnBOzaV_4gM5noPX8YCb_p06d8PJvGfFy3Y,118
 bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
 bayinx/dists/gamma2.py,sha256=MuFudL2UTfk8HgWVofNaR36JTmUpmtxvg1Mifu98MvM,1567
@@ -29,7 +29,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.7.dist-info/METADATA,sha256=bQGouAjty73m1UeFCOWgRMw7Is0ffja7xDXAyS-EzDM,3079
-bayinx-0.3.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.3.7.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
-bayinx-0.3.7.dist-info/RECORD,,
+bayinx-0.3.9.dist-info/METADATA,sha256=yn3AC0J8_xLnLYEEocDZyNRlyyuN2Ys8C--i9AFFYeY,3079
+bayinx-0.3.9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.3.9.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
+bayinx-0.3.9.dist-info/RECORD,,