bayinx 0.3.10__py3-none-any.whl → 0.5.3__py3-none-any.whl

bayinx/mhx/vi/normalizing_flow.py

@@ -1,149 +0,0 @@
-from typing import Any, Generic, Self, Tuple, TypeVar
-
-import equinox as eqx
-import jax.flatten_util as jfu
-import jax.numpy as jnp
-import jax.random as jr
-import jax.tree_util as jtu
-from jaxtyping import Array, Key, Scalar
-
-from bayinx.core import Flow, Model, Variational
-
-M = TypeVar('M', bound=Model)
-class NormalizingFlow(Variational, Generic[M]):
-    """
-    An ordered collection of diffeomorphisms that map a base distribution to a
-    normalized approximation of a posterior distribution.
-
-    # Attributes
-    - `base`: A base variational distribution.
-    - `flows`: An ordered collection of continuously parameterized diffeomorphisms.
-    """
-
-    flows: list[Flow]
-    base: Variational
-
-    def __init__(self, base: Variational, flows: list[Flow], model: M):
-        """
-        Constructs an unoptimized normalizing flow posterior approximation.
-
-        # Parameters
-        - `base`: The base variational distribution.
-        - `flows`: A list of diffeomorphisms.
-        - `model`: A probabilistic `Model` object.
-        """
-        # Partition model
-        params, self._constraints = eqx.partition(model, model.filter_spec)
-
-        # Flatten params component
-        _, self._unflatten = jfu.ravel_pytree(params)
-
-        self.base = base
-        self.flows = flows
-
-    @property
-    @eqx.filter_jit
-    def filter_spec(self):
-        # Generate empty specification
-        filter_spec = jtu.tree_map(lambda _: False, self)
-
-        # Specify variational parameters based on each flow's filter spec.
-        filter_spec = eqx.tree_at(
-            lambda vari: vari.flows,
-            filter_spec,
-            replace=[flow.filter_spec for flow in self.flows],
-        )
-
-        return filter_spec
-
-    @eqx.filter_jit
-    def sample(self, n: int, key: Key = jr.PRNGKey(0)):
-        """
-        Sample from the variational distribution `n` times.
-        """
-        # Sample from the base distribution
-        draws: Array = self.base.sample(n, key)
-
-        # Apply forward transformations
-        for map in self.flows:
-            draws = map.forward(draws)
-
-        return draws
-
-    @eqx.filter_jit
-    def eval(self, draws: Array) -> Array:
-        # Evaluate base density
-        variational_evals: Array = self.base.eval(draws)
-
-        for map in self.flows:
-            # Compute adjustment
-            draws, laj = map.adjust_density(draws)
-
-            # Adjust variational density
-            variational_evals = variational_evals - laj
-
-        return variational_evals
-
-    @eqx.filter_jit
-    def __eval(self, draws: Array, data=None) -> Tuple[Array, Array]:
-        """
-        Evaluate the posterior and variational densities together at the
-        transformed `draws` to avoid extra compute.
-
-        # Parameters
-        - `draws`: Draws from the base variational distribution.
-        - `data`: Any data required to evaluate the posterior density.
-
-        # Returns
-        The posterior and variational densities as JAX Arrays.
-        """
-        # Evaluate base density
-        variational_evals: Array = self.base.eval(draws)
-
-        for map in self.flows:
-            # Compute adjustment
-            draws, laj = map.adjust_density(draws)
-
-            # Adjust variational density
-            variational_evals = variational_evals - laj
-
-        # Evaluate posterior at final variational draws
-        posterior_evals = self.eval_model(draws, data)
-
-        return posterior_evals, variational_evals
-
-    @eqx.filter_jit
-    def elbo(self, n: int, key: Key = jr.PRNGKey(0), data: Any = None) -> Scalar:
-        dyn, static = eqx.partition(self, self.filter_spec)
-
-        @eqx.filter_jit
-        def elbo(dyn: Self, n: int, key: Key, data: Any = None):
-            self = eqx.combine(dyn, static)
-
-            # Sample draws from variational distribution
-            draws: Array = self.base.sample(n, key)
-
-            posterior_evals, variational_evals = self.__eval(draws, data)
-            # Evaluate ELBO
-            return jnp.mean(posterior_evals - variational_evals)
-
-        return elbo(dyn, n, key, data)
-
-    @eqx.filter_jit
-    def elbo_grad(self, n: int, key: Key, data: Any = None) -> Self:
-        dyn, static = eqx.partition(self, self.filter_spec)
-
-        @eqx.filter_grad
-        @eqx.filter_jit
-        def elbo_grad(dyn: Self, n: int, key: Key, data: Any = None):
-            self = eqx.combine(dyn, static)
-
-            # Sample draws from variational distribution
-            draws: Array = self.base.sample(n, key)
-
-            posterior_evals, variational_evals = self.__eval(draws, data)
-
-            # Evaluate ELBO
-            return jnp.mean(posterior_evals - variational_evals)
-
-        return elbo_grad(dyn, n, key, data)

bayinx/mhx/vi/standard.py

@@ -1,63 +0,0 @@
-
-import equinox as eqx
-import jax.numpy as jnp
-import jax.random as jr
-import jax.tree_util as jtu
-from jax.flatten_util import ravel_pytree
-from jaxtyping import Array, Key
-
-from bayinx.core._variational import M, Variational
-from bayinx.dists import normal
-
-
-class Standard(Variational[M]):
-    """
-    A standard normal approximation to a posterior distribution.
-
-    # Attributes
-    - `dim`: Dimension of the parameter space.
-    """
-    dim: int
-
-    def __init__(self, model: M):
-        """
-        Constructs a standard normal approximation to a posterior distribution.
-
-        # Parameters
-        - `model`: A probabilistic `Model` object.
-        """
-        # Partition model
-        params, self._constraints = eqx.partition(model, model.filter_spec)
-
-        # Flatten params component
-        params, self._unflatten = ravel_pytree(params)
-
-        # Store dimension of parameter space
-        self.dim = jnp.size(params)
-
-    @eqx.filter_jit
-    def sample(self, n: int, key: Key = jr.PRNGKey(0)) -> Array:
-        # Sample variational draws
-        draws: Array = jr.normal(key=key, shape=(n, self.dim))
-
-        return draws
-
-    @eqx.filter_jit
-    def eval(self, draws: Array) -> Array:
-        return normal.logprob(
-            x=draws,
-            mu=jnp.array(0.0),
-            sigma=jnp.array(1.0),
-        ).sum(axis=1, keepdims=True)
-
-    @property
-    def filter_spec(self):
-        filter_spec = jtu.tree_map(lambda _: False, self)
-
-        return filter_spec
-
-    def elbo(self):
-        return None
-
-    def elbo_grad(self):
-        return None

bayinx-0.3.10.dist-info/METADATA

@@ -1,39 +0,0 @@
-Metadata-Version: 2.4
-Name: bayinx
-Version: 0.3.10
-Summary: Bayesian Inference with JAX
-License-File: LICENSE
-Requires-Python: >=3.12
-Requires-Dist: equinox>=0.11.12
-Requires-Dist: jax>=0.4.38
-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>
-
-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`).
-
-In the short-term, I'm going to focus on:
-1) Implementing as much machinery as I feel is enough.
-2) Figuring out how to design the `Model` superclass to have something like the `transformed pars {}` block but unifies transformations and constraints.
-3) Figuring out how to design the library to automatically recognize what kind of machinery is amenable to a given probabilistic model.
-
-In the long-term, I'm going to focus on:
-1) How to get `Stan`-like declarative syntax in Python with minimal syntactic overhead(to get as close as possible to statements like `X ~ Normal(mu, 1)`), while also allowing users to work with `target` directly when needed(same as `Stan` does).
-2) How to make working with the posterior as easy as possible.
-    - That's a vague goal but practically it means how to easily evaluate statements like $P(\theta \in [-1, 1] | \mathcal{D}, \mathcal{M})$, or set up contrasts and evaluate $P(\mu_1 - \mu_2 > 0 | \mathcal{D}, \mathcal{M})$, or simulate the posterior predictive to generate plots, etc.
-
-Although this is somewhat separate from the goals of the project, if this does pan out how I'm invisioning it I'd like an R formula-like syntax to shorten model construction in scenarios where the model is just a GLMM or similar(think `brms`).
-
-Additionally, when I get around to it I'd like the package documentation to also include theoretical and implementation details for all machinery implemented(with overthinking boxes because I do like that design from McElreath's book).
-
-
-# TODO
-- Find some way to discern between models with all floating-point parameters and weirder models with integer parameters. Useful for restricting variational methods like `MeanField` to `Model`s that only have floating-point parameters.
-- Look into adaptively tuning ADAM hyperparameters.
-- Control variates for meanfield VI? Look at https://proceedings.mlr.press/v33/ranganath14.html more closely.
-- Low-rank affine flow?
-- https://arxiv.org/pdf/1803.05649 implement sylvester flows.
-- Learn how to generate documentation lol.
-- Figure out how to make transform_pars for flows such that there is no performance loss. Noticing some weird behaviour when adding constraints.

bayinx-0.3.10.dist-info/RECORD

@@ -1,35 +0,0 @@
-bayinx/__init__.py,sha256=TM-aoRaPX6jSYtCM7Jv59TPV-H6bcDk1-VMttYP1KME,99
-bayinx/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bayinx/constraints/__init__.py,sha256=PiWXZKi7YdbTMKvw-OE5f-t87jJT893uAFrwWWBfOdg,64
-bayinx/constraints/lower.py,sha256=30y0l6PF-tbS9LR_tto9AvwmsvXq1ExU-v8DLrJD4g4,1446
-bayinx/core/__init__.py,sha256=bZvQITgW0DWuPKl3wCLKt6WHKogYKx8Zz36g8z9Aung,253
-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/_parameter.py,sha256=r20JedTW2lY0miNNh9y6LeIVAsGX1kP_rlGxphW_jZg,1080
-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
-bayinx/dists/normal.py,sha256=Yc2X8F7JoLYwprtK8bA2BPva1tAY7MEs3oSk5pMortI,3822
-bayinx/dists/posnormal.py,sha256=w9plA1EctXwXOiY0doc4ZndjnwptbEZBHHCGdc4gviY,7292
-bayinx/dists/uniform.py,sha256=7XgVvOrzINEFA6HJTYUOFwlWhEtrQQQ1aPJ_ZLOzLEc,2365
-bayinx/dists/censored/__init__.py,sha256=UVihMbQgAzCoOk_Zt5wrumPv5-acuTzV3TYMB-U1gOc,49
-bayinx/dists/censored/gamma2/__init__.py,sha256=GO3jIF1En0ZxYF5JqvC0helLAL6yv8-LG6Ih2NOUYQc,33
-bayinx/dists/censored/gamma2/r.py,sha256=dKAOYstufwgDwibQZHrJxA1d2gawj-7K3IkaCRCzNTg,2446
-bayinx/dists/censored/posnormal/__init__.py,sha256=GO3jIF1En0ZxYF5JqvC0helLAL6yv8-LG6Ih2NOUYQc,33
-bayinx/dists/censored/posnormal/r.py,sha256=Ypi6w_t53pAzRVzjcStx2RhozkAlCDLnQmgKykhpQQ4,3426
-bayinx/mhx/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-bayinx/mhx/vi/__init__.py,sha256=2woNB5oZxfs8pZCkOfzriGahRFLzkLdkTj8_keTN0I0,205
-bayinx/mhx/vi/meanfield.py,sha256=Z7kGQAyp5iB8rEdjbwAbVTFH4GwxlTKDZFbdJ-FN5Vs,3739
-bayinx/mhx/vi/normalizing_flow.py,sha256=8pLMDdZPIt5wlgbhHWSFY1ChSWM9pvSD2bQx3zgz1F8,4710
-bayinx/mhx/vi/standard.py,sha256=W-ZvigJkUpqVlREgiFm9io8ansT1XpZwq5AqSmdv--E,1578
-bayinx/mhx/vi/flows/__init__.py,sha256=Hn0Wqvvyv8Vr-mFmimwgNKCByxj-fjrlIvdR7tUSolg,180
-bayinx/mhx/vi/flows/fullaffine.py,sha256=11y_A0oO3bkKDSz-EQ6Sf4Ec2M7vHZxw94EdvADwVYQ,1954
-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.10.dist-info/METADATA,sha256=y99n8rNP62ezyDKJSmBQvUQx52WKUEgyHtw3QnnVrvs,3080
-bayinx-0.3.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.3.10.dist-info/licenses/LICENSE,sha256=VMhLhj5hx6VAENZBaNfXrmsNl7ov9uRh0jZ6D3ltgv4,1070
-bayinx-0.3.10.dist-info/RECORD,,