bayinx 0.1.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

bayinx/__init__.py

@@ -0,0 +1 @@
+from bayinx.core.model import Model as Model

bayinx/core/__init__.py

@@ -0,0 +1,3 @@
+from bayinx.core.flow import Flow as Flow
+from bayinx.core.model import Model as Model
+from bayinx.core.variational import Variational as Variational

bayinx/core/flow.py

@@ -0,0 +1,68 @@
+from abc import abstractmethod
+from typing import Callable, Dict, Self, Tuple
+
+import equinox as eqx
+from jaxtyping import Array, Float
+
+from bayinx.core.utils import __MyMeta
+
+
+class Flow(eqx.Module, metaclass=__MyMeta):
+    """
+    A superclass used to define continuously parameterized diffeomorphisms for normalizing flows.
+
+    # Attributes
+    - `pars`: A dictionary of JAX Arrays representing parameters of the diffeomorphism.
+    - `constraints`: A dictionary of functions that constrain their corresponding parameter.
+    """
+
+    params: Dict[str, Float[Array, "..."]]
+    constraints: Dict[str, Callable[[Float[Array, "..."]], Float[Array, "..."]]]
+
+    @abstractmethod
+    def forward(self, draws: Array) -> Array:
+        """
+        Computes the forward transformation of `draws`.
+        """
+        pass
+
+    @abstractmethod
+    def adjust_density(self, draws: Array) -> Tuple[Array, Array]:
+        """
+        Computes the log-absolute-determinant of the Jacobian at `draws` and applies the forward transformation.
+
+        # Returns
+        A tuple of JAX Arrays containing the log-absolute-determinant of the Jacobians and transformed draws.
+        """
+        pass
+
+    def __init_subclass__(cls):
+        # Add contrain_pars method
+        def constrain_pars(self: Self):
+            """
+            Constrain `params` to the appropriate domain.
+
+            # Returns
+            A dictionary of transformed JAX Arrays representing the constrained parameters.
+            """
+            t_params = self.params
+
+            for par, map in self.constraints.items():
+                t_params[par] = map(t_params[par])
+
+            return t_params
+
+        cls.constrain_pars = eqx.filter_jit(constrain_pars)
+
+        # Add transform_pars method if not present
+        if not callable(getattr(cls, "transform_pars", None)):
+            def transform_pars(self: Self) -> Dict[str, Array]:
+                """
+                Apply a custom transformation to `params` if needed.
+
+                # Returns
+                A dictionary of transformed JAX Arrays representing the transformed parameters.
+                """
+                return self.constrain_pars()
+
+            cls.transform_pars = eqx.filter_jit(transform_pars)

bayinx/core/model.py

@@ -0,0 +1,55 @@
+from abc import abstractmethod
+from typing import Any, Callable, Dict
+
+import equinox as eqx
+from jaxtyping import Array, Scalar
+
+from bayinx.core.utils import __MyMeta
+
+
+class Model(eqx.Module, metaclass=__MyMeta):
+    """
+    A superclass used to define probabilistic models.
+
+    # Attributes
+    - `params`: A dictionary of JAX Arrays representing parameters of the model.
+    - `constraints`: A dictionary of functions that constrain their corresponding parameter.
+    """
+
+    params: Dict[str, Array]
+    constraints: Dict[str, Callable[[Array], Array]]
+
+    @abstractmethod
+    def eval(self, data: Any) -> Scalar:
+        pass
+
+    def __init_subclass__(cls):
+        # Add constrain method
+        def constrain_pars(self: Model) -> Dict[str, Array]:
+            """
+            Constrain `params` to the appropriate domain.
+
+            # Returns
+            A dictionary of transformed JAX Arrays representing the constrained parameters.
+            """
+            t_params = self.params
+
+            for par, map in self.constraints.items():
+                t_params[par] = map(t_params[par])
+
+            return t_params
+
+        cls.constrain_pars = eqx.filter_jit(constrain_pars)
+
+        # Add transform_pars method if not present
+        if not callable(getattr(cls, "transform_pars", None)):
+            def transform_pars(self: Model) -> Dict[str, Array]:
+                """
+                Apply a custom transformation to `params` if needed.
+
+                # Returns
+                A dictionary of transformed JAX Arrays representing the transformed parameters.
+                """
+                return self.constrain_pars()
+
+            cls.transform_pars = eqx.filter_jit(transform_pars)

bayinx/core/utils.py

@@ -0,0 +1,54 @@
+from typing import Callable, Dict
+
+import equinox as eqx
+from jaxtyping import Array
+
+
+class __MyMeta(type(eqx.Module)):
+    """
+    Metaclass to ensure attribute types are respected.
+    """
+
+    def __call__(cls, *args, **kwargs):
+        obj = super().__call__(*args, **kwargs)
+
+        # Check parameters are a Dict of JAX Arrays
+        if not isinstance(obj.params, Dict):
+            raise ValueError(
+                f"Model {cls.__name__} must initialize 'params' as a dictionary."
+            )
+
+        for key, value in obj.params.items():
+            if not isinstance(value, Array):
+                raise TypeError(f"Parameter '{key}' must be a JAX Array.")
+
+        # Check constraints are a Dict of functions
+        if not isinstance(obj.constraints, Dict):
+            raise ValueError(
+                f"Model {cls.__name__} must initialize 'constraints' as a dictionary."
+            )
+
+        for key, value in obj.constraints.items():
+            if not isinstance(value, Callable):
+                raise TypeError(f"Constraint for parameter '{key}' must be a function.")
+
+        # Check that the constrain method returns a dict equivalent to `params`
+        t_params: Dict[str, Array] = obj.constrain_pars()
+
+        if not isinstance(t_params, Dict):
+            raise ValueError(
+                f"The 'constrain' method of {cls.__name__} must return a Dict of JAX Arrays."
+            )
+
+        for key, value in t_params.items():
+            if not isinstance(value, Array):
+                raise TypeError(f"Constrained parameter '{key}' must be a JAX Array.")
+
+            if not value.shape == obj.params[key].shape:
+                raise ValueError(
+                    f"Constrained parameter '{key}' must have same shape as unconstrained counterpart."
+                )
+
+        # Check transform_pars
+
+        return obj

bayinx/core/variational.py

@@ -0,0 +1,159 @@
+from abc import abstractmethod
+from typing import Any, Callable, Self, Tuple
+
+import equinox as eqx
+import jax
+import jax.lax as lax
+import jax.numpy as jnp
+import jax.random as jr
+import optax as opx
+from jaxtyping import Array, Float, Key, PyTree, Scalar
+from optax import GradientTransformation, OptState, Schedule
+
+from bayinx.core import Model
+
+
+class Variational(eqx.Module):
+    """
+    A superclass used to define variational methods.
+
+    # Attributes
+    - `_unflatten`: A static function to transform draws from the variational distribution back to a `Model`.
+    - `_constraints`: A static partitioned `Model` with the constraints of the `Model` used to initialize the `Variational` object.
+    """
+
+    _unflatten: Callable[[Float[Array, "..."]], Model]
+    _constraints: Model
+
+    @abstractmethod
+    def sample(self, n: int, key: Key) -> Array:
+        """
+        Sample from the variational distribution.
+        """
+        pass
+
+    @abstractmethod
+    def eval(self, draws: Array) -> Array:
+        """
+        Evaluate the variational distribution at `draws`.
+        """
+        pass
+
+    @abstractmethod
+    def elbo(self, n: int, key: Key, data: Any = None) -> Array:
+        """
+        Evaluate the ELBO.
+        """
+        pass
+
+    @abstractmethod
+    def elbo_grad(self, n: int, key: Key, data: Any = None) -> PyTree:
+        """
+        Evaluate the gradient of the ELBO.
+        """
+        pass
+
+    @abstractmethod
+    def filter_spec(self):
+        """
+        Filter specification for dynamic and static components of the `Variational`.
+        """
+        pass
+
+    def __init_subclass__(cls):
+        """
+        Construct methods that are shared across all VI methods.
+        """
+
+        def eval_model(self, draws: Array, data: Any = None) -> Array:
+            """
+            Reconstruct models from variational draws and evaluate their posterior density.
+
+            # Parameters
+            - `draws`: A set of variational draws.
+            - `data`: Data used to evaluate the posterior(if needed).
+            """
+            # Unflatten variational draw
+            model: Model = self._unflatten(draws)
+
+            # Combine with constraints
+            model: Model = eqx.combine(model, self._constraints)
+
+            # Evaluate posterior density
+            return model.eval(data)
+
+        cls.eval_model = jax.vmap(eqx.filter_jit(eval_model), (None, 0, None))
+
+        def fit(
+            self,
+            max_iters: int,
+            data: Any = None,
+            learning_rate: float = 1,
+            tolerance: float = 1e-4,
+            var_draws: int = 1,
+            key: Key = jr.PRNGKey(0),
+        ) -> Self:
+            """
+            Optimize the variational distribution.
+
+            # Parameters
+            - `max_iters`: Maximum number of iterations for the optimization loop.
+            - `data`: Data to evaluate the posterior density with(if available).
+            - `learning_rate`: Initial learning rate for optimization.
+            - `tolerance`: Relative tolerance of ELBO decrease for stopping early.
+            - `var_draws`: Number of variational draws to draw each iteration.
+            - `key`: A PRNG key.
+            """
+            # Construct scheduler
+            schedule: Schedule = opx.cosine_decay_schedule(
+                init_value=learning_rate, decay_steps=max_iters
+            )
+
+            # Initialize optimizer
+            optim: GradientTransformation = opx.chain(
+                opx.scale(-1.0), opx.nadam(schedule)
+            )
+            opt_state: OptState = optim.init(eqx.filter(self, self.filter_spec()))
+
+            # Optimization loop helper functions
+            @eqx.filter_jit
+            def condition(state: Tuple[Self, OptState, Scalar, Key]):
+                # Unpack iteration state
+                self, opt_state, i, key = state
+
+                return i < max_iters
+
+            @eqx.filter_jit
+            def body(state: Tuple[Self, OptState, Scalar, Key]):
+                # Unpack iteration state
+                self, opt_state, i, key = state
+
+                # Update iteration
+                i = i + 1
+
+                # Update PRNG key
+                key, _ = jr.split(key)
+
+                # Compute gradient of the ELBO
+                updates: PyTree = self.elbo_grad(var_draws, key, data)
+
+                # Compute updates
+                updates, opt_state = optim.update(
+                    updates, opt_state, eqx.filter(self, self.filter_spec())
+                )
+
+                # Update variational distribution
+                self: Self = eqx.apply_updates(self, updates)
+
+                return self, opt_state, i, key
+
+            # Run optimization loop
+            self = lax.while_loop(
+                cond_fun=condition,
+                body_fun=body,
+                init_val=(self, opt_state, jnp.array(0, jnp.uint32), key),
+            )[0]
+
+            return self
+
+        cls.fit = eqx.filter_jit(fit)

bayinx/dists/bernoulli.py

@@ -0,0 +1,33 @@
+import jax.lax as lax
+from jaxtyping import Array, ArrayLike, Real, UInt
+
+
+# MARK: Functions ----
+def prob(x: UInt[ArrayLike, "..."], p: Real[ArrayLike, "..."]) -> Real[Array, "..."]:
+    """
+    The probability mass function (PMF) for a Bernoulli distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the PDF.
+    - `p`:     The probability parameter(s).
+
+    # Returns
+    The PMF evaluated at `x`. The output will have the broadcasted shapes of `x` and `p`.
+    """
+
+    return lax.pow(p, x) * lax.pow(1 - p, 1 - x)
+
+
+def logprob(x: UInt[ArrayLike, "..."], p: Real[ArrayLike, "..."]) -> Real[Array, "..."]:
+    """
+    The log probability mass function (log PMF) for a Bernoulli distribution.
+
+    # Parameters
+    - `x`:  Value(s) at which to evaluate the log PMF.
+    - `p`:  The probability parameter(s).
+
+    # Returns
+    The log PMF evaluated at `x`. The output will have the broadcasted shapes of `x` and `p`.
+    """
+
+    return lax.log(p) * x + (1 - x) * lax.log(1 - p)

bayinx/dists/normal.py

@@ -0,0 +1,83 @@
+# MARK: Imports ----
+import jax.lax as _lax
+
+## Typing
+from jaxtyping import Array, Real
+
+# MARK: Constants
+_PI = 3.141592653589793
+
+
+# MARK: Functions ----
+def prob(
+    x: Real[Array, "..."], mu: Real[Array, "..."], sigma: Real[Array, "..."]
+) -> Real[Array, "..."]:
+    """
+    The probability density function (PDF) for a Normal distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the PDF.
+    - `mu`:     The mean/location parameter(s).
+    - `sigma`:  The non-negative standard deviation parameter(s).
+
+    # Returns
+    The PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    """
+
+    return _lax.exp(-0.5 * _lax.square((x - mu) / sigma)) / (
+        sigma * _lax.sqrt(2.0 * _PI)
+    )
+
+
+def logprob(
+    x: Real[Array, "..."], mu: Real[Array, "..."], sigma: Real[Array, "..."]
+) -> Real[Array, "..."]:
+    """
+    The log of the probability density function (log PDF) for a Normal distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the log PDF.
+    - `mu`:     The mean/location parameter(s).
+    - `sigma`:  The non-negative standard deviation parameter(s).
+
+    # Returns
+    The log of the PDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    """
+
+    return -_lax.log(sigma * _lax.sqrt(2.0 * _PI)) - 0.5 * _lax.square((x - mu) / sigma)
+
+
+def uprob(
+    x: Real[Array, "..."], mu: Real[Array, "..."], sigma: Real[Array, "..."]
+) -> Real[Array, "..."]:
+    """
+    The unnormalized probability density function (uPDF) for a Normal distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the uPDF.
+    - `mu`:     The mean/location parameter(s).
+    - `sigma`:  The non-negative standard deviation parameter(s).
+
+    # Returns
+    The uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    """
+
+    return _lax.exp(-0.5 * _lax.square((x - mu) / sigma)) / sigma
+
+
+def ulogprob(
+    x: Real[Array, "..."], mu: Real[Array, "..."], sigma: Real[Array, "..."]
+) -> Real[Array, "..."]:
+    """
+    The log of the unnormalized probability density function (log uPDF) for a Normal distribution.
+
+    # Parameters
+    - `x`:      Value(s) at which to evaluate the log uPDF.
+    - `mu`:     The mean/location parameter(s).
+    - `sigma`:  The non-negative standard deviation parameter(s).
+
+    # Returns
+    The log uPDF evaluated at `x`. The output will have the broadcasted shapes of `x`, `mu`, and `sigma`.
+    """
+
+    return -_lax.log(sigma) - 0.5 * _lax.square((x - mu) / sigma)

bayinx/machinery/variational/__init__.py

@@ -0,0 +1,5 @@
+from bayinx.machinery.variational.meanfield import MeanField as MeanField
+from bayinx.machinery.variational.normalizing_flow import (
+    NormalizingFlow as NormalizingFlow,
+)
+from bayinx.machinery.variational.standard import Standard as Standard

bayinx/machinery/variational/flows/__init__.py

@@ -0,0 +1,3 @@
+from bayinx.machinery.variational.flows.affine import Affine as Affine
+from bayinx.machinery.variational.flows.planar import Planar as Planar
+from bayinx.machinery.variational.flows.radial import Radial as Radial

bayinx/machinery/variational/flows/affine.py

@@ -0,0 +1,68 @@
+from functools import partial
+from typing import Callable, Dict, Tuple
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array, Float, Scalar
+
+from bayinx.core import Flow
+
+
+class Affine(Flow):
+    """
+    An affine flow.
+
+    # Attributes
+    - `params`: A dictionary containing the JAX Arrays representing the scale and shift parameters.
+    - `constraints`: A dictionary of constraining transformations.
+    """
+
+    params: Dict[str, Float[Array, "..."]]
+    constraints: Dict[str, Callable[[Float[Array, "..."]], Float[Array, "..."]]] = (
+        eqx.field(static=True)
+    )
+
+    def __init__(self, dim: int):
+        """
+        Initializes an affine flow.
+
+        # Parameters
+        - `dim`: The dimension of the parameter space.
+        """
+        self.params = {
+            "shift": jnp.zeros(dim),
+            "scale": jnp.zeros((dim, dim)),
+        }
+
+        self.constraints = {"scale": lambda m: jnp.tril(jnp.exp(m))}
+
+    @eqx.filter_jit
+    def forward(self, draws: Array) -> Array:
+        params = self.constrain_pars()
+
+        # Extract parameters
+        shift: Array = params["shift"]
+        scale: Array = params["scale"]
+
+        # Compute forward transformation
+        draws = draws @ scale + shift
+
+        return draws
+
+    @eqx.filter_jit
+    @partial(jax.vmap, in_axes=(None, 0))
+    def adjust_density(self, draws: Array) -> Tuple[Scalar, Array]:
+        params = self.constrain_pars()
+
+        # Extract parameters
+        shift: Array = params["shift"]
+        scale: Array = params["scale"]
+
+        # Compute forward transformation
+        draws = draws @ scale + shift
+
+        # Compute ladj
+        ladj: Scalar = jnp.log(jnp.diag(scale)).sum()
+
+        return ladj, draws

bayinx/machinery/variational/flows/planar.py

@@ -0,0 +1,76 @@
+from functools import partial
+from typing import Callable, Dict, Tuple
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jax.random as jr
+from jaxtyping import Array, Float, Scalar
+
+from bayinx.core import Flow
+
+
+class Planar(Flow):
+    """
+    A planar flow.
+
+    # Attributes
+    - `params`: A dictionary containing the JAX Arrays representing the flow parameters.
+    - `constraints`: A dictionary of constraining transformations.
+    """
+
+    params: Dict[str, Float[Array, "..."]]
+    constraints: Dict[str, Callable[[Float[Array, "..."]], Float[Array, "..."]]] = (
+        eqx.field(static=True)
+    )
+
+    def __init__(self, dim: int, key=jr.PRNGKey(0)):
+        """
+        Initializes a planar flow.
+
+        # Parameters
+        - `dim`: The dimension of the parameter space.
+        """
+        self.params = {
+            "u": jr.normal(key, (dim,)),
+            "w": jr.normal(key, (dim,)),
+            "b": jr.normal(key, (1,)),
+        }
+        self.constraints = {}
+
+    @eqx.filter_jit
+    @partial(jax.vmap, in_axes=(None, 0))
+    def forward(self, draws: Array) -> Array:
+        params = self.constrain_pars()
+
+        # Extract parameters
+        w: Array = params["w"]
+        u: Array = params["u"]
+        b: Array = params["b"]
+
+        # Compute forward transformation
+        draws = draws + u * jnp.tanh(draws.dot(w) + b)
+
+        return draws
+
+    @eqx.filter_jit
+    @partial(jax.vmap, in_axes=(None, 0))
+    def adjust_density(self, draws: Array) -> Tuple[Scalar, Array]:
+        params = self.constrain_pars()
+
+        # Extract parameters
+        w: Array = params["w"]
+        u: Array = params["u"]
+        b: Array = params["b"]
+
+        # Compute shared intermediates
+        x: Array = draws.dot(w) + b
+
+        # Compute forward transformation
+        draws = draws + u * jnp.tanh(x)
+
+        # Compute ladj
+        h_prime: Scalar = 1.0 - jnp.square(jnp.tanh(x))
+        ladj: Scalar = jnp.log(jnp.abs(1.0 + h_prime * u.dot(w)))
+
+        return ladj, draws

bayinx/machinery/variational/flows/radial.py

@@ -0,0 +1,95 @@
+from functools import partial
+from typing import Callable, Dict, Tuple
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jax.random as jr
+from jax.numpy.linalg import norm
+from jaxtyping import Array, Float, Scalar
+
+from bayinx.core import Flow
+
+
+class Radial(Flow):
+    """
+    A radial flow.
+
+    # Attributes
+    - `params`: A dictionary containing the JAX Arrays representing the flow parameters.
+    - `constraints`: A dictionary of constraining transformations.
+    """
+
+    params: Dict[str, Float[Array, "..."]]
+    constraints: Dict[str, Callable[[Float[Array, "..."]], Float[Array, "..."]]] = (
+        eqx.field(static=True)
+    )
+
+    def __init__(self, dim: int, key=jr.PRNGKey(0)):
+        """
+        Initializes a planar flow.
+
+        # Parameters
+        - `dim`: The dimension of the parameter space.
+        """
+        self.params = {
+            "alpha": jnp.array(1.0),
+            "beta": jnp.array(1.0),
+            "center": jnp.ones(dim),
+        }
+        self.constraints = {"beta": jnp.exp}
+
+    @partial(jax.vmap, in_axes=(None, 0))
+    @eqx.filter_jit
+    def forward(self, draws: Array) -> Array:
+        """
+        Applies the forward radial transformation for each draw.
+
+        # Parameters
+        - `draws`: Draws from some layer of a normalizing flow.
+
+        # Returns
+        The transformed samples.
+        """
+        params = self.transform_pars()
+
+        # Extract parameters
+        alpha = params["alpha"]
+        beta = params["beta"]
+        center = params["center"]
+
+        # Compute distance to center per-draw
+        r: Array = norm(draws - center)
+
+        # Apply forward transformation
+        draws = draws + (beta / (alpha + r)) * (draws - center)
+
+        return draws
+
+    @partial(jax.vmap, in_axes=(None, 0))
+    @eqx.filter_jit
+    def adjust_density(self, draws: Array) -> Tuple[Scalar, Array]:
+        params = self.transform_pars()
+
+        # Extract parameters
+        alpha = params["alpha"]
+        beta = params["beta"]
+        center = params["center"]
+
+        # Compute distance to center per-draw
+        r: Array = norm(draws - center)
+
+        # Compute shared intermediates
+        x: Array = beta / (alpha + r)
+
+        # Apply forward transformation
+        draws = draws + (x) * (draws - center)
+
+        # Compute density adjustment
+        ladj = jnp.log(
+            jnp.abs(
+                (1.0 + alpha * beta / (alpha + r) ** 2.0) * (1.0 + x) ** (center.size - 1.0)
+            )
+        )
+
+        return ladj, draws

bayinx/machinery/variational/flows/sylvester.py

@@ -0,0 +1,76 @@
+from functools import partial
+from typing import Callable, Dict, Tuple
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jax.random as jr
+from jaxtyping import Array, Float, Scalar
+
+from bayinx.core import Flow
+
+
+class Sylvester(Flow):
+    """
+    A sylvester flow.
+
+    # Attributes
+    - `params`: A dictionary containing the JAX Arrays representing the flow parameters.
+    - `constraints`: A dictionary of constraining transformations.
+    """
+
+    params: Dict[str, Float[Array, "..."]]
+    constraints: Dict[str, Callable[[Float[Array, "..."]], Float[Array, "..."]]] = (
+        eqx.field(static=True)
+    )
+
+    def __init__(self, dim: int, key=jr.PRNGKey(0)):
+        """
+        Initializes a planar flow.
+
+        # Parameters
+        - `dim`: The dimension of the parameter space.
+        """
+        self.params = {
+            "u": jr.normal(key, (dim,)),
+            "w": jr.normal(key, (dim,)),
+            "b": jr.normal(key, (1,)),
+        }
+        self.constraints = {}
+
+    @eqx.filter_jit
+    @partial(jax.vmap, in_axes=(None, 0))
+    def forward(self, draws: Array) -> Array:
+        params = self.constrain_pars()
+
+        # Extract parameters
+        w: Array = params["w"]
+        u: Array = params["u"]
+        b: Array = params["b"]
+
+        # Compute forward transformation
+        draws = draws + u * jnp.tanh(draws.dot(w) + b)
+
+        return draws
+
+    @eqx.filter_jit
+    @partial(jax.vmap, in_axes=(None, 0))
+    def adjust_density(self, draws: Array) -> Tuple[Scalar, Array]:
+        params = self.constrain_pars()
+
+        # Extract parameters
+        w: Array = params["w"]
+        u: Array = params["u"]
+        b: Array = params["b"]
+
+        # Compute shared intermediates
+        x: Array = draws.dot(w) + b
+
+        # Compute forward transformation
+        draws = draws + u * jnp.tanh(x)
+
+        # Compute ladj
+        h_prime: Scalar = 1.0 - jnp.square(jnp.tanh(x))
+        ladj: Scalar = jnp.log(jnp.abs(1.0 + h_prime * u.dot(w)))
+
+        return ladj, draws

bayinx/machinery/variational/meanfield.py

@@ -0,0 +1,124 @@
+from typing import Any, Callable, Dict, Self
+
+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, Float, Key, Scalar
+
+from bayinx.core import Model, Variational
+from bayinx.dists import normal
+
+
+class MeanField(Variational):
+    """
+    A fully factorized Gaussian approximation to a posterior distribution.
+
+    # Attributes
+    - `var_params`: The variational parameters for the approximation.
+    """
+
+    var_params: Dict[str, Float[Array, "..."]]
+    _unflatten: Callable[[Float[Array, "..."]], Model] = eqx.field(static=True)
+    _constraints: Model = eqx.field(static=True)
+
+    def __init__(self, model: Model):
+        """
+        Constructs an unoptimized meanfield posterior approximation.
+
+        # Parameters
+        - `model`: A probabilistic `Model` object.
+        """
+        # Partition model
+        params, self._constraints = eqx.partition(model, eqx.is_array)
+
+        # Flatten params component
+        flat_params, self._unflatten = ravel_pytree(params)
+
+        # Initialize variational parameters
+        self.var_params = {
+            "mean": flat_params,
+            "log_std": jnp.zeros(flat_params.size, dtype=flat_params.dtype),
+        }
+
+    @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.var_params["mean"].size))
+            * jnp.exp(self.var_params["log_std"])
+            + self.var_params["mean"]
+        )
+
+        return draws
+
+    @eqx.filter_jit
+    def eval(self, draws: Array) -> Array:
+        return normal.logprob(
+            x=draws,
+            mu=self.var_params["mean"],
+            sigma=jnp.exp(self.var_params["log_std"]),
+        ).sum(axis=1)
+
+    @eqx.filter_jit
+    def filter_spec(self):
+        filter_spec = jtu.tree_map(lambda _: False, self)
+        filter_spec = eqx.tree_at(
+            lambda mf: mf.var_params,
+            filter_spec,
+            replace=True,
+        )
+        return filter_spec
+
+    @eqx.filter_jit
+    def elbo(self, n: int, key: Key, data: Any = None) -> Scalar:
+        """
+        Estimate the ELBO and its gradient(w.r.t the variational parameters).
+        """
+        # Partition variational
+        dyn, static = eqx.partition(self, self.filter_spec())
+
+        @eqx.filter_jit
+        def elbo(dyn: Self, n: int, key: Key, data: Any = None) -> Scalar:
+            # Combine
+            vari = eqx.combine(dyn, static)
+
+            # Sample draws from variational distribution
+            draws: Array = vari.sample(n, key)
+
+            # Evaluate posterior density for each draw
+            posterior_evals: Array = vari.eval_model(draws, data)
+
+            # Evaluate variational density for each draw
+            variational_evals: Array = vari.eval(draws)
+
+            # 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:
+        # Partition
+        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):
+            # Combine
+            vari = eqx.combine(dyn, static)
+
+            # Sample draws from variational distribution
+            draws: Array = vari.sample(n, key)
+
+            # Evaluate posterior density for each draw
+            posterior_evals: Array = vari.eval_model(draws, data)
+
+            # Evaluate variational density for each draw
+            variational_evals: Array = vari.eval(draws)
+
+            # Evaluate ELBO
+            return jnp.mean(posterior_evals - variational_evals)
+
+        return elbo_grad(dyn, n, key, data)

bayinx/machinery/variational/normalizing_flow.py

@@ -0,0 +1,152 @@
+from typing import Any, Callable, Self, Tuple
+
+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, Float, Key, Scalar
+
+from bayinx.core import Flow, Model, Variational
+
+
+class NormalizingFlow(Variational):
+    """
+    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
+    _unflatten: Callable[[Float[Array, "..."]], Model] = eqx.field(static=True)
+    _constraints: Model = eqx.field(static=True)
+
+    def __init__(self, base: Variational, flows: list[Flow], model: Model):
+        """
+        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, eqx.is_array)
+
+        # Flatten params component
+        flat_params, self._unflatten = jfu.ravel_pytree(params)
+
+        self.base = base
+        self.flows = flows
+
+    @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
+            ladj, draws = map.adjust_density(draws)
+
+            # Adjust variational density
+            variational_evals = variational_evals - ladj
+
+        return variational_evals
+
+    @eqx.filter_jit
+    def _eval(self, draws: Array, data=None) -> Tuple[Scalar, Array]:
+        """
+        Evaluate the posterior and variational densities at the transformed
+        `draws` to avoid extra compute when requiring variational draws for
+        the posterior evaluation.
+
+        # Parameters
+        - `draws`: Draws from the base variational distribution.
+        - `data`: Any data required to evaluate the posterior density.
+
+        # Returns
+        The posterior and variational densities.
+        """
+        # Evaluate base density
+        variational_evals: Array = self.base.eval(draws)
+
+        for map in self.flows:
+            # Compute adjustment
+            ladj, draws = map.adjust_density(draws)
+
+            # Adjust variational density
+            variational_evals = variational_evals - ladj
+
+        # Evaluate posterior at final variational draws
+        posterior_evals = self.eval_model(draws, data)
+
+        return posterior_evals, variational_evals
+
+    def filter_spec(self):
+        # Only optimize the parameters of the flows
+        filter_spec = jtu.tree_map(lambda _: False, self)
+        filter_spec = eqx.tree_at(
+            lambda nf: nf.flows,
+            filter_spec,
+            replace=True,
+        )
+
+        return filter_spec
+
+    @eqx.filter_jit
+    def elbo(self, n: int, key: Key, data: Any = None) -> Scalar:
+        # Partition
+        dyn, static = eqx.partition(self, self.filter_spec())
+
+        @eqx.filter_jit
+        def elbo(dyn: Self, n: int, key: Key, data: Any = None):
+            # Combine
+            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:
+        # Partition
+        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):
+            # Combine
+            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/machinery/variational/standard.py

@@ -0,0 +1,67 @@
+from typing import Callable
+
+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, Float, Key
+
+from bayinx.core import Model, Variational
+from bayinx.dists import normal
+
+
+class Standard(Variational):
+    """
+    A standard normal distribution approximation to a posterior distribution.
+
+    # Attributes
+    - `dim`: Dimension of the parameter space.
+    """
+
+    dim: int = eqx.field(static=True)
+    _unflatten: Callable[[Float[Array, "..."]], Model] = eqx.field(static=True)
+    _constraints: Model = eqx.field(static=True)
+
+    def __init__(self, model: Model):
+        """
+        Constructs a standard normal approximation to a posterior distribution.
+
+        # Parameters
+        - `model`: A probabilistic `Model` object.
+        """
+        # Partition model
+        _, self._constraints = eqx.partition(model, eqx.is_array)
+
+        # Flatten params component
+        _, self._unflatten = ravel_pytree(_)
+
+        # Store dimension of parameter space
+        self.dim = jnp.size(_)
+
+    @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)
+
+    @eqx.filter_jit
+    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.2.2.dist-info/METADATA

@@ -0,0 +1,40 @@
+Metadata-Version: 2.4
+Name: bayinx
+Version: 0.2.2
+Summary: Bayesian Inference with JAX
+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
+Requires-Dist: pytest-benchmark>=5.1.0
+Requires-Dist: pytest>=8.3.5
+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{M})$, or set up contrasts and evaluate $P(\mu_1 - \mu_2 > 0 | \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.2.2.dist-info/RECORD

@@ -0,0 +1,26 @@
+bayinx/__init__.py,sha256=l20JdkSsE_XGZlZFNEtySXf4NIlbjrao14vXPB-H6aQ,45
+bayinx/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/core/__init__.py,sha256=7vW2F8t3K4TWlSu5nZrYCdUrz5N9FMIfQQBn3IoeH6o,150
+bayinx/core/flow.py,sha256=4vj1t2xNPGp1VPE4xUshY-rHAw__KvSwjGDtKkW2taE,2252
+bayinx/core/model.py,sha256=AI4eHrXAds3K7eWgZ9g5E6Kh76HP6WTn6s6q_0tnhck,1719
+bayinx/core/utils.py,sha256=-YewhqzMFL3GJEjVdm3LgaZyHwDs9IVYllU9wAXZrtw,1859
+bayinx/core/variational.py,sha256=T42uUNkF2tP1HJPyeIv7ISdika_G28wR_OOFXzx_hgo,4978
+bayinx/dists/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/dists/bernoulli.py,sha256=xMV9BgtVX_1XkPdZ43q0meMIEkgMyuUPx--dyo6_DKs,1006
+bayinx/dists/binomial.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/dists/gamma.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/dists/gamma2.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/dists/normal.py,sha256=e9gXXAHeZQKjBndW2TnMvP3gtmvpfYGG7kehcpGeAoU,2590
+bayinx/machinery/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bayinx/machinery/variational/__init__.py,sha256=5GdhqBOHKsXg2tZGAMNlxyrLPD0-s64wAEy8998cHZ4,247
+bayinx/machinery/variational/meanfield.py,sha256=stI96DNhHNIROnr8rLNEN9SN_0lXqkwit0KNto34Q6A,3889
+bayinx/machinery/variational/normalizing_flow.py,sha256=qypgPq9vIqSIJNOHCDaN-hvwFfttNQ5_yXqvmi5hslI,4796
+bayinx/machinery/variational/standard.py,sha256=IQdNd5QIE8u3zcOw7K4EW69lIQ0ZUGGDvwZVyvrYHxA,1739
+bayinx/machinery/variational/flows/__init__.py,sha256=VGh-ffuUfMso_0JxwGCJQ2yVnFJdOrkFsSnorojQldY,213
+bayinx/machinery/variational/flows/affine.py,sha256=TPyUUPRoSkyDMwGO5wtq-Ei8DAUvlb_N6JCk7uPlbJQ,1748
+bayinx/machinery/variational/flows/planar.py,sha256=rJ1XpqoWzig_5Udq6oCh5JV4ptlTRLRS7tb9DCX22lE,2013
+bayinx/machinery/variational/flows/radial.py,sha256=NF1tCd_PH6m8eqjJkom2c30sRUQ04Vf8zeRx_RQCDcg,2526
+bayinx/machinery/variational/flows/sylvester.py,sha256=DeZl4Fkz9XpCGsfDcjS0eWlrMR0xMO0MPfJvoDhixSA,2019
+bayinx-0.2.2.dist-info/METADATA,sha256=gNOkUv-EdtqnLp154L29ycWd4t2c2wsRmyqhNXYc1C8,3099
+bayinx-0.2.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bayinx-0.2.2.dist-info/RECORD,,

bayinx-0.1.0.dist-info/METADATA

@@ -1,8 +0,0 @@
-Metadata-Version: 2.4
-Name: bayinx
-Version: 0.1.0
-Summary: A personal library for Bayesian inference
-Requires-Python: >=3.12
-Requires-Dist: baycompx
-Requires-Dist: jax>=0.4.38
-Requires-Dist: jaxtyping>=0.2.36

bayinx-0.1.0.dist-info/RECORD

@@ -1,5 +0,0 @@
-bayinx/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bayinx/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bayinx-0.1.0.dist-info/METADATA,sha256=vo6K-1pFW6eW-SeCGyfJSxjkJ92h21qxUgLeVCL8C3Q,209
-bayinx-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bayinx-0.1.0.dist-info/RECORD,,