bayinx 0.3.6__py3-none-any.whl → 0.3.8__py3-none-any.whl

bayinx/core/_variational.py

@@ -106,8 +106,8 @@ class Variational(eqx.Module, Generic[M]):
         dyn, static = eqx.partition(self, self.filter_spec)
 
         # Construct scheduler
-        schedule: Schedule = opx.cosine_decay_schedule(
-            init_value=learning_rate, decay_steps=max_iters
+        schedule: Schedule = opx.warmup_cosine_decay_schedule(
+            init_value=1e-16, peak_value=learning_rate, warmup_steps=int(max_iters/10), decay_steps=max_iters-int(max_iters/10)
         )
 
         # Initialize optimizer
@@ -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], 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,7 +171,7 @@ 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)
 

{bayinx-0.3.6.dist-info → bayinx-0.3.8.dist-info}/METADATA

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

{bayinx-0.3.6.dist-info → bayinx-0.3.8.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=szm1WuUh_3pxzFfQy92TR4p2Sk-fR6rO-4-LrJMeVGI,5356
+bayinx/core/_variational.py,sha256=omZ3jFrYNdjfDVAOmNhjPtOrJcBD6yFlypvFROn6GgU,5467
 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.6.dist-info/METADATA,sha256=WEdMVyISWGgK0KJvuSlkpbObsxiVfGvIxky7OsuYdXg,3079
-bayinx-0.3.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.3.6.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
-bayinx-0.3.6.dist-info/RECORD,,
+bayinx-0.3.8.dist-info/METADATA,sha256=Z6oJm_2TgfJyhle7aG-DeDiW3LqE6Mcwt8iNku5ai54,3079
+bayinx-0.3.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.3.8.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
+bayinx-0.3.8.dist-info/RECORD,,