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

bayinx/core/_optimization.py

@@ -1,4 +1,4 @@
-from typing import Any, Tuple, TypeVar
+from typing import Any, Callable, Tuple, TypeVar
 
 import equinox as eqx
 import jax.lax as lax
@@ -20,7 +20,7 @@ def optimize_model(
     tolerance: float = 1e-4,
 ) -> M:
     """
-    Optimize the dynamic parameters of the model.
+    Optimize the parameters of the model.
 
     # Parameters
     - `max_iters`: Maximum number of iterations for the optimization loop.
@@ -33,14 +33,13 @@ def optimize_model(
     dyn, static = eqx.partition(model, model.filter_spec)
 
     # Derive gradient for posterior
-    @eqx.filter_jit
-    @eqx.filter_grad
-    def eval_grad(dyn: M):
+    def eval(dyn: M) -> Scalar:
         # Reconstruct model
         model: M = eqx.combine(dyn, static)
 
         # Evaluate posterior
         return model.eval(data)
+    eval_grad: Callable[[M], M] = eqx.filter_jit(eqx.filter_grad(eval))
 
     # Construct scheduler
     schedule: Schedule = opx.warmup_cosine_decay_schedule(
@@ -71,7 +70,7 @@ def optimize_model(
         i = i + 1
 
         # Evaluate gradient of posterior
-        updates = eval_grad(dyn)
+        updates: PyTree = eval_grad(dyn)
 
         # Compute updates
         updates, opt_state = optim.update(

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/dists/gamma2.py

@@ -46,4 +46,4 @@ def logprob(
         + nu * (lax.log(nu) - lax.log(mu))
         + (nu - 1.0) * lax.log(x)
         - (x * nu / mu)
-    )  # pyright: ignore
+    )

bayinx/dists/posnormal.py

@@ -139,9 +139,6 @@ def cdf(
 
     # Returns
     The CDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
-
-    # Notes
-    Not numerically stable for small `x`.
     """
     # Cast to Array
     x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)
@@ -160,7 +157,6 @@ def cdf(
     return evals
 
 
-# TODO: make numerically stable
 def logcdf(
     x: Float[ArrayLike, "..."],
     mu: Float[ArrayLike, "..."],
@@ -176,9 +172,6 @@ def logcdf(
 
     # Returns
     The log CDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
-
-    # Notes
-    Not numerically stable for small `x`.
     """
     # Cast to Array
     x, mu, sigma = jnp.asarray(x), jnp.asarray(mu), jnp.asarray(sigma)

bayinx/mhx/vi/meanfield.py

@@ -23,12 +23,13 @@ class MeanField(Variational, Generic[M]):
 
     var_params: Dict[str, Float[Array, "..."]]  # todo: just expand to attributes
 
-    def __init__(self, model: M):
+    def __init__(self, model: M, init_log_std: float = 0.0):
         """
         Constructs an unoptimized meanfield posterior approximation.
 
         # Parameters
         - `model`: A probabilistic `Model` object.
+        - `init_log_std`: The initial log-transformed standard deviation of the Gaussian approximation.
         """
         # Partition model
         params, self._constraints = eqx.partition(model, model.filter_spec)
@@ -39,7 +40,7 @@ class MeanField(Variational, Generic[M]):
         # Initialize variational parameters
         self.var_params = {
             "mean": params,
-            "log_std": jnp.zeros(params.size, dtype=params.dtype),
+            "log_std": jnp.full(params.size, init_log_std, params.dtype),
         }
 
     @property

{bayinx-0.3.14.dist-info → bayinx-0.3.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bayinx
-Version: 0.3.14
+Version: 0.3.16
 Summary: Bayesian Inference with JAX
 License-File: LICENSE
 Requires-Python: >=3.12
@@ -10,9 +10,9 @@ Requires-Dist: jaxtyping>=0.2.36
 Requires-Dist: optax>=0.2.4
 Description-Content-Type: text/markdown
 
-# <ins>Bay</ins>esian <ins>In</ins>ference with JA<ins>X</ins>
+# `Bayinx`: <ins>Bay</ins>esian <ins>In</ins>ference with JA<ins>X</ins>
 
-The endgoal of this project is to build a Bayesian inference library that is similar in feel to `Stan`(where you can define a probabilistic model with syntax that is equivalent to how you would write it out on a chalkboard) but allows for arbitrary models(e.g., ones with discrete parameters) and offers a suite of "machinery" to fit the model; this means I want to expand upon `Stan`'s existing toolbox of methods for estimation(point optimization, variational methods, MCMC) while keeping everything performant(hence using `JAX`).
+The endgoal of this project is to build a Bayesian inference library that is similar in feel to `Stan`(where you can define a probabilistic model with syntax that is similar to how you would write it out on a chalkboard) but allows for arbitrary models(e.g., ones with discrete parameters) and offers a suite of "machinery" to fit the model; this means I want to expand upon `Stan`'s existing toolbox of methods for estimation(point optimization, variational methods, MCMC) while keeping everything performant(hence using `JAX`).
 
 In the short-term, I'm going to focus on:
 1) Implementing as much machinery as I feel is enough.

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

@@ -6,14 +6,14 @@ bayinx/core/__init__.py,sha256=Qmy0EjzqqKwI9F8rjmC9j6J8hiDw6A54yOck2WuQJkY,344
 bayinx/core/_constraint.py,sha256=Gx07ZT66VE2y-qZCmBDm3_y0wO4xQyslZW10Lec1_lM,761
 bayinx/core/_flow.py,sha256=3q4rKvATnbUpuj4ICUd4lIxu_3z7GRDuNujVhAku1X0,2342
 bayinx/core/_model.py,sha256=FJUyYVE9e2uTFamxtSMKY_VV2stiU2QF68Wl_7EAKEU,2895
-bayinx/core/_optimization.py,sha256=dL3COMFTP0_FeD44hreNybh_UD6zqJy6sSP54ITJsBc,2605
+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=MuFudL2UTfk8HgWVofNaR36JTmUpmtxvg1Mifu98MvM,1567
+bayinx/dists/gamma2.py,sha256=HtB60LUQdPj4yDAHme2jsHNmLfrAKWsSZnDYkxAGaOI,1548
 bayinx/dists/normal.py,sha256=Yc2X8F7JoLYwprtK8bA2BPva1tAY7MEs3oSk5pMortI,3822
-bayinx/dists/posnormal.py,sha256=w9plA1EctXwXOiY0doc4ZndjnwptbEZBHHCGdc4gviY,7292
+bayinx/dists/posnormal.py,sha256=cOLCdd39DX3v8DD-seSIKNk4OfdNfaYaLzpCh_xBGyw,7150
 bayinx/dists/uniform.py,sha256=2ZQxEfAX5TFgSPuQ8joFDuFbd_NfmQ1GvmGGjusqvNQ,3461
 bayinx/dists/censored/__init__.py,sha256=UVihMbQgAzCoOk_Zt5wrumPv5-acuTzV3TYMB-U1gOc,49
 bayinx/dists/censored/gamma2/__init__.py,sha256=GO3jIF1En0ZxYF5JqvC0helLAL6yv8-LG6Ih2NOUYQc,33
@@ -23,7 +23,7 @@ bayinx/dists/censored/posnormal/r.py,sha256=wMDt2Am1TD376ms8B-o6PFCJZXmUJd2-aBC-
 bayinx/mhx/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 bayinx/mhx/opt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bayinx/mhx/vi/__init__.py,sha256=3T1dEpiiRge4tW-vpS0xBob_RbO1iVFnL3fVCRUawCM,205
-bayinx/mhx/vi/meanfield.py,sha256=kx5WeD93-XO8NHxd4L4pZ8V19Y9B6j-yE3Y5OBXMcTk,3743
+bayinx/mhx/vi/meanfield.py,sha256=wM0v6Q2m0pPyEdOaT8DvFqdRLsHijr7AYNQKAcsZtRQ,3881
 bayinx/mhx/vi/normalizing_flow.py,sha256=vzLu5H1G1-pBqhgHWmIZkUTyPE1DxC9vBwpiZeIyu1I,4712
 bayinx/mhx/vi/standard.py,sha256=LYgglaGQMGmXpzFR4eMJnXkl2PhBeggbXMvO5zJpf2c,1578
 bayinx/mhx/vi/flows/__init__.py,sha256=Hn0Wqvvyv8Vr-mFmimwgNKCByxj-fjrlIvdR7tUSolg,180
@@ -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.14.dist-info/METADATA,sha256=ba_0a1aYSrKIGhGFG_IWA7TltEmGpK78IqCFeLV7XVI,3080
-bayinx-0.3.14.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.3.14.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
-bayinx-0.3.14.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,,